diff --git a/.dockerignore b/.dockerignore index 51781b9463e80602979ea1acc12d8383a0f3cfdd..d17ca3b6a12c7f56c70438e0f3819458a7fbebad 100644 --- a/.dockerignore +++ b/.dockerignore @@ -13,6 +13,4 @@ notebooks/ .idea/ .vscode/ .DS_Store -# Exclude most JSON files but allow agents_config.json *.json -!agents_config.json diff --git a/.env-template b/.env-template index 6f7b3903fbb6e3b8bc9833c05849de2e20f09a12..9c0eb6184347aca72ab9a5d75c533772dcb9fe8a 100644 --- a/.env-template +++ b/.env-template @@ -7,9 +7,10 @@ GROQ_API_KEY=your-groq-api-key-here ANTHROPIC_API_KEY=your-anthropic-api-key-here GEMINI_API_KEY=your-gemini-api-key-here -ADMIN_API_KEY=admin123 +DATABASE_URL=postgresql://dbadmin:admin123@auto-analyst-db.xxxxxxxxxxxxx.us-east-1.rds.amazonaws.com:5432/autoanalyst +DB_HOST=auto-analyst-db.xxxxxxxxxxxxx.us-east-1.rds.amazonaws.com +DB_NAME=autoanalyst +DB_USER=dbadmin +DB_PASSWORD=admin123 -DATABASE_URL=sqlite:///chat_database.db -ENVIRONMENT="development" - -FRONTEND_URL="http://localhost:3000/" +ENV="development" \ No newline at end of file diff --git a/.gitattributes b/.gitattributes index 2da0042e9d73d886f7617b53307b1fa1250bf413..cf41ff0aa3d485b555adb8d32a5dff3f6ae85b41 100644 --- a/.gitattributes +++ b/.gitattributes @@ -35,4 +35,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text images/Auto-Analyst[[:space:]]Banner.png filter=lfs diff=lfs merge=lfs -text chat_database.db filter=lfs diff=lfs merge=lfs -text -*.png filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index ca1f7136632753c9c4096c013216d3c992808be3..e7e1d31c5dc90bb4ec460c449086d3085d30bded 100644 --- a/.gitignore +++ b/.gitignore @@ -11,12 +11,14 @@ venv/ try* +chat_database_temp.db + logs/ updated_code.py sample_code.py -*.duckdb + *.dump migrations/ @@ -24,29 +26,3 @@ migrations/ *.pyc alembic.ini - -*.db - -schema*.md - -# agent_config.json - - -notebooks/ - -*.xlsx -*.xls -*.xlsm -*.xlsb - - -testing.ipynb -redis_index.json -email_to_userid_mapping.json -redis_backup_20250906_143859.json -"*.db" -"*.sqlite" -"*.sqlite3" -"venv/" -"__pycache__/" -"*.pyc" diff --git a/.huggingface.yaml b/.huggingface.yaml deleted file mode 100644 index bb327d84f7b0af37f4674dde2d21d4e05dae301a..0000000000000000000000000000000000000000 --- a/.huggingface.yaml +++ /dev/null @@ -1 +0,0 @@ -sdk: docker diff --git a/Dockerfile b/Dockerfile index 2aed07ea03091304a0ff6f58cf8fbe7c3662e053..480e44cadc9d966416c932e3b3bec93394022595 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,29 +6,8 @@ ENV PATH="/home/user/.local/bin:$PATH" WORKDIR /app - - COPY --chown=user ./requirements.txt requirements.txt RUN pip install --no-cache-dir --upgrade -r requirements.txt COPY --chown=user . /app - -# Verify agents_config.json was copied (it should be in the backend directory) -RUN if [ -f "/app/agents_config.json" ]; then \ - echo "✅ agents_config.json found in container"; \ - ls -la /app/agents_config.json; \ - else \ - echo "⚠️ agents_config.json not found in container - will use fallback templates"; \ - ls -la /app/ | grep -E "agents|config" || echo "No config files found"; \ - fi - -# Make entrypoint script executable -USER root -RUN chmod +x /app/entrypoint_local.sh -# Make populate script executable -RUN chmod +x /app/scripts/populate_agent_templates.py - -USER user - -# Use the entrypoint script instead of directly running uvicorn -CMD ["/app/entrypoint_local.sh"] \ No newline at end of file +CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"] \ No newline at end of file diff --git a/Procfile b/Procfile deleted file mode 100644 index ba1ce3a2ec5dd7aefe3fe32c76143e59ad526b3d..0000000000000000000000000000000000000000 --- a/Procfile +++ /dev/null @@ -1 +0,0 @@ -web: (python scripts/init_production_db.py || echo "DB init failed") && (python scripts/populate_agent_templates.py || echo "Template init failed") && uvicorn app:app --host 0.0.0.0 --port $PORT diff --git a/agents_config.json b/agents_config.json deleted file mode 100644 index 8bb72b3586ffeaa246c9fad8b0e01d431dcc31ac..0000000000000000000000000000000000000000 --- a/agents_config.json +++ /dev/null @@ -1,147 +0,0 @@ -{ - "templates": [ - { - "template_name": "preprocessing_agent", - "display_name": "Data Preprocessing Agent", - "description": "Cleans and prepares a DataFrame using Pandas and NumPy—handles missing values, detects column types, and converts date strings to datetime", - "icon_url": "/icons/templates/preprocessing_agent.svg", - "category": "Data Manipulation", - "is_premium_only": false, - "variant_type": "individual", - "base_agent": "preprocessing_agent", - "is_active": true, - "prompt_template": "You are a AI data-preprocessing agent. Generate clean and efficient Python code using NumPy and Pandas to perform introductory data preprocessing on a pre-loaded DataFrame df, based on the user's analysis goals.\nPreprocessing Requirements:\n1. Identify Column Types\n- Separate columns into numeric and categorical using:\n categorical_columns = df.select_dtypes(include=[object, 'category']).columns.tolist()\n numeric_columns = df.select_dtypes(include=[np.number]).columns.tolist()\n2. Handle Missing Values\n- Numeric columns: Impute missing values using the mean of each column\n- Categorical columns: Impute missing values using the mode of each column\n3. Convert Date Strings to Datetime\n- For any column suspected to represent dates (in string format), convert it to datetime using:\n def safe_to_datetime(date):\n try:\n return pd.to_datetime(date, errors='coerce', cache=False)\n except (ValueError, TypeError):\n return pd.NaT\n df['datetime_column'] = df['datetime_column'].apply(safe_to_datetime)\n- Replace 'datetime_column' with the actual column names containing date-like strings\nImportant Notes:\n- Do NOT create a correlation matrix — correlation analysis is outside the scope of preprocessing\n- Do NOT generate any plots or visualizations\nOutput Instructions:\n1. Include the full preprocessing Python code\n2. Provide a brief bullet-point summary of the steps performed. Example:\n• Identified 5 numeric and 4 categorical columns\n• Filled missing numeric values with column means\n• Filled missing categorical values with column modes\n• Converted 1 date column to datetime format\n Respond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "planner_preprocessing_agent", - "display_name": "Data Preprocessing Agent", - "description": "Multi-agent planner variant: Cleans and prepares a DataFrame using Pandas and NumPy—handles missing values, detects column types, and converts date strings to datetime", - "icon_url": "/icons/templates/preprocessing_agent.svg", - "category": "Data Manipulation", - "is_premium_only": false, - "variant_type": "planner", - "base_agent": "preprocessing_agent", - "is_active": true, - "prompt_template": "You are a data preprocessing agent optimized for multi-agent data analytics pipelines.\n\nYou are given:\n* A raw dataset (often just uploaded or loaded).\n* A user-defined goal (e.g., clean data for analysis, prepare for modeling).\n***plan_instructions** containing:\n ***'create'**: Variables you must create (e.g., ['df_cleaned', 'preprocessing_summary', 'column_types'])\n ***'use'**: Variables you must use (e.g., ['df', 'raw_data'])\n * **'instruction'**: Specific preprocessing instructions\n\n### Your Planner-Optimized Responsibilities:\n* **ALWAYS follow plan_instructions** - essential for pipeline data flow\n* Create ONLY the variables specified in plan_instructions['create']\n* Use ONLY the variables specified in plan_instructions['use']\n* Apply preprocessing as per plan_instructions['instruction']\n* Ensure cleaned data integrates seamlessly with downstream agents\n\n### Core Preprocessing Techniques:\n* Identify and categorize column types (numeric, categorical, datetime)\n* Handle missing values appropriately:\n - Numeric: impute with mean, median, or specified strategy\n - Categorical: impute with mode or specified strategy\n* Convert date strings to datetime format with proper error handling\n* Remove duplicates and handle data quality issues\n* Apply data type optimizations for memory efficiency\n* Create preprocessing summaries for pipeline transparency\n\n### Multi-Agent Best Practices:\n* Use exact variable names from plan_instructions['create']\n* Ensure data format compatibility for downstream agents\n* Maintain data integrity and schema consistency\n* Document preprocessing steps for pipeline reproducibility\n\n### Output:\n* Python code implementing preprocessing per plan_instructions\n* Summary of data cleaning and transformation operations\n* Focus on seamless integration with analysis and modeling agents\n\nRespond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "statistical_analytics_agent", - "display_name": "Statistical Analytics Agent", - "description": "Performs statistical analysis (e.g., regression, seasonal decomposition) using statsmodels, with proper handling of categorical data and missing values", - "icon_url": "/icons/templates/statsmodel.svg", - "category": "Data Modelling", - "is_premium_only": false, - "variant_type": "individual", - "base_agent": "statistical_analytics_agent", - "is_active": true, - "prompt_template": "You are a statistical analytics agent. Your task is to take a dataset and a user-defined goal and output Python code that performs the appropriate statistical analysis to achieve that goal. Follow these guidelines:\nIMPORTANT: You may be provided with previous interaction history. The section marked \"### Current Query:\" contains the user's current request. Any text in \"### Previous Interaction History:\" is for context only and is NOT part of the current request.\nData Handling:\nAlways handle strings as categorical variables in a regression using statsmodels C(string_column).\nDo not change the index of the DataFrame.\nConvert X and y into float when fitting a model.\nError Handling:\nAlways check for missing values and handle them appropriately.\nEnsure that categorical variables are correctly processed.\nProvide clear error messages if the model fitting fails.\nRegression:\nFor regression, use statsmodels and ensure that a constant term is added to the predictor using sm.add_constant(X).\nHandle categorical variables using C(column_name) in the model formula.\nFit the model with model = sm.OLS(y.astype(float), X.astype(float)).fit().\nSeasonal Decomposition:\nEnsure the period is set correctly when performing seasonal decomposition.\nVerify the number of observations works for the decomposition.\nOutput:\nEnsure the code is executable and as intended.\nAlso choose the correct type of model for the problem\nAvoid adding data visualization code.\nProvide a concise bullet-point summary of the statistical analysis performed.\n\nExample Summary:\n• Applied linear regression with OLS to predict house prices based on 5 features\n• Model achieved R-squared of 0.78\n• Significant predictors include square footage (p<0.001) and number of bathrooms (p<0.01)\n• Detected strong seasonal pattern with 12-month periodicity\n• Forecast shows 15% growth trend over next quarter\nRespond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "planner_statistical_analytics_agent", - "display_name": "Statistical Analytics Agent", - "description": "Multi-agent planner variant: Performs statistical analysis (e.g., regression, seasonal decomposition) using statsmodels, with proper handling of categorical data and missing values", - "icon_url": "/icons/templates/statsmodel.svg", - "category": "Data Modelling", - "is_premium_only": false, - "variant_type": "planner", - "base_agent": "statistical_analytics_agent", - "is_active": true, - "prompt_template": "You are tasked with performing statistical analysis on datasets based on provided structured inputs. Ensure comprehensive results by following these detailed instructions carefully:\n\n### Input Format:\nYou will receive structured input, which includes:\n1. **Dataset Description**:\n - Overview of the dataset, including its purpose and key columns (types, etc.).\n - Specific preprocessing instructions for each column, particularly for data type conversions and missing value handling.\n\n2. **Analytical Goal**:\n - A clearly defined goal, such as generating specific insights, performing calculations, or summarizing the data.\n\n3. **Plan Instructions**:\n - Detailed actions that should be taken, outlining what variables to create, which existing variables to use, and any other necessary processing steps.\n\n### Key Responsibilities:\n1. **Data Preprocessing**:\n - Inspect columns for needed preprocessing according to the dataset description provided.\n - Implement preprocessing as specified, including handling categorical variables with appropriate encoding (e.g., one-hot encoding).\n\n2. **Statistical Analysis**:\n - Conduct analysis based on the defined goal, which may involve:\n - Descriptive statistics (means, medians, etc.).\n - Correlation analysis to understand relationships among numerical variables.\n - Calculation of specific metrics described in the task.\n - Utilize libraries such as `pandas` for data manipulation and `numpy` for numerical operations.\n\n3. **Output**:\n - Results must be presented in a structured and organized text format, integrating all specified variables into the final report.\n - Avoid creating any intermediates that are not specified in the plan instructions.\n\n4. **Error Handling**:\n - Integrate error checks to confirm that all requisite variables are well defined and valid prior to executing operations.\n - Address edge cases, including situations where DataFrames may be empty or lack the necessary columns.\n\n5. **Documentation**:\n - Summarize all findings succinctly, detailing:\n - Key statistical outcomes, highlighting identifiable trends or relationships.\n - Potential data quality issues, such as missing values or outliers.\n\n### Analytical Methodology:\n- Always start with data cleaning, ensuring that missing values are handled as specified (e.g., filling with mean or median) and outlier checks are sufficient.\n- When performing statistical analysis, use measures that facilitate understanding of data distributions, such as means, medians, and standard deviations, as well as categorizations based on quantitative thresholds.\n- Implement segmentation strategies based on calculated scores, specify the thresholds clearly for different segments, and ensure that insights can lead to actionable outcomes.\n- Include plots where required, and ensure they are prepared in a separate stage, if indicated in the plan.\n\n### Important Notes:\n- Do not modify data indexes unless instructed; maintain the integrity of the dataset structure throughout.\n- Ensure all numerical data is converted to the appropriate types prior to analysis.\n- In the event that visualizations are indicated, prepare these in a separate task as per the capabilities outlined.\n\nBy adhering to these instructions meticulously, you will deliver consistent and high-quality analytical insights tailored to the provided datasets." }, - { - "template_name": "data_viz_agent", - "display_name": "Data Visualization Agent", - "description": "Creates interactive data visualizations using Plotly with advanced styling and formatting options", - "icon_url": "/icons/templates/plotly.svg", - "category": "Data Visualization", - "is_premium_only": false, - "variant_type": "individual", - "base_agent": "data_viz_agent", - "is_active": true, - "prompt_template": "You are an AI agent responsible for generating interactive data visualizations using Plotly.\nIMPORTANT Instructions:\n- The section marked \"### Current Query:\" contains the user's request. Any text in \"### Previous Interaction History:\" is for context only and should NOT be treated as part of the current request.\n- You must only use the tools provided to you. This agent handles visualization only.\n- If len(df) > 50000, always sample the dataset before visualization using: \nif len(df) > 50000: \n df = df.sample(50000, random_state=1)\n- Each visualization must be generated as a **separate figure** using go.Figure(). \nDo NOT use subplots under any circumstances.\n- Each figure must be returned individually using: \nfig.to_html(full_html=False)\n- Use update_layout with xaxis and yaxis **only once per figure**.\n- Enhance readability and clarity by: \n• Using low opacity (0.4-0.7) where appropriate \n• Applying visually distinct colors for different elements or categories \n- Make sure the visual **answers the user's specific goal**: \n• Identify what insight or comparison the user is trying to achieve \n• Choose the visualization type and features (e.g., color, size, grouping) to emphasize that goal \n• For example, if the user asks for \"trends in revenue,\" use a time series line chart; if they ask for \"top-performing categories,\" use a bar chart sorted by value \n• Prioritize highlighting patterns, outliers, or comparisons relevant to the question\n- Never include the dataset or styling index in the output.\n- If there are no relevant columns for the requested visualization, respond with: \n\"No relevant columns found to generate this visualization.\"\n- Use only one number format consistently: either 'K', 'M', or comma-separated values like 1,000/1,000,000. Do not mix formats.\n- Only include trendlines in scatter plots if the user explicitly asks for them.\n- Output only the code and a concise bullet-point summary of what the visualization reveals.\n- Always end each visualization with: \nfig.to_html(full_html=False)\nRespond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "sk_learn_agent", - "display_name": "Machine Learning Agent", - "description": "Trains and evaluates machine learning models using scikit-learn, including classification, regression, and clustering with feature importance insights", - "icon_url": "/icons/templates/sk_learn_agent.svg", - "category": "Data Modelling", - "is_premium_only": false, - "variant_type": "individual", - "base_agent": "sk_learn_agent", - "is_active": true, - "prompt_template": "You are a machine learning agent. \nYour task is to take a dataset and a user-defined goal, and output Python code that performs the appropriate machine learning analysis to achieve that goal. \nYou should use the scikit-learn library.\nIMPORTANT: You may be provided with previous interaction history. The section marked \"### Current Query:\" contains the user's current request. Any text in \"### Previous Interaction History:\" is for context only and is NOT part of the current request.\nMake sure your output is as intended!\nProvide a concise bullet-point summary of the machine learning operations performed.\n\nExample Summary:\n• Trained a Random Forest classifier on customer churn data with 80/20 train-test split\n• Model achieved 92% accuracy and 88% F1-score\n• Feature importance analysis revealed that contract length and monthly charges are the strongest predictors of churn\n• Implemented K-means clustering (k=4) on customer shopping behaviors\n• Identified distinct segments: high-value frequent shoppers (22%), occasional big spenders (35%), budget-conscious regulars (28%), and rare visitors (15%)\nRespond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "planner_data_viz_agent", - "display_name": "Data Visualization Agent", - "description": "Multi-agent planner variant: Creates interactive data visualizations using Plotly with advanced styling and formatting options", - "icon_url": "/icons/templates/plotly.svg", - "category": "Data Visualization", - "is_premium_only": false, - "variant_type": "planner", - "base_agent": "data_viz_agent", - "is_active": true, - "prompt_template": "You are a data visualization agent designed to generate effective visualizations based on user-defined goals and specific datasets provided in a structured format. Your enhanced responsibilities and necessary details for best practices are as follows:\n\n### Input Format:\n1. **Dataset**: Provided in JSON or Pandas DataFrame format, detailing its structure and attributes, including column types, preprocessing requirements, and guidelines on handling missing values.\n\n2. **Goal**: A clear statement that defines the analytical objectives for visualization (e.g., performance analysis, relationship discovery, or data clustering).\n\n3. **Plan Instructions**: Specific directives from an analytical planner regarding analysis creation, dataset usage, and additional plotting notes.\n\n4. **Styling Index**: Contains visual preferences for the plots, axis specifications, formatting requirements, and any template references.\n\n### Responsibilities:\n1. **Data Handling**:\n - Confirm the presence of necessary data variables before proceeding.\n - If datasets exceed 50,000 rows, sample them down to 5,000 rows for efficiency.\n - Check for missing values in crucial columns and address them according to preprocessing instructions (e.g., mean or median imputation).\n - Ensure that columns have consistent lengths, especially those involved in calculations.\n\n2. **Visualization Creation**:\n - Utilize Plotly and Matplotlib for visualization, focusing on user-defined goals and creation instructions from the plan.\n - Generate multiple relevant visualizations based on specific goals, potentially including bar charts, histograms, scatter plots, word clouds, or heatmaps as dictated by the task requirements.\n - Implement text processing techniques for natural language data (e.g., removing special characters while preserving language integrity).\n - For datasets comprising categorical variables, ensure they are handled correctly, including appropriate encoding of categorical features and filling in missing data with default categories.\n\n3. **Layout and Styling**:\n - Follow the provided styling index for clarity and aesthetics, ensuring cohesive axis formatting and color usage.\n - Use `update_yaxes` and `update_xaxes` for effective axis presentation, maintaining a uniform look across visualizations.\n\n4. **Error Handling**:\n - If essential variables are missing or if there are mismatched array lengths, return clear error messages indicating the specific issues (e.g., DataFrame not defined Column missing).\n - Address any ambiguities in input formats and expectations proactively rather than making unfounded assumptions.\n\n5. **Output**:\n - Visualizations must be displayed using the appropriate methods such as `.show()` or `fig.to_html(full_html=False)` for seamless HTML rendering.\n - Each visualization should include comprehensive legends or annotations where applicable, helping to clarify complex data stories.\n\n### Domain-Specific Considerations:\n- **Text Data**: When handling natural language data, particularly in non-English languages, use regular expressions to efficiently clean and preprocess text while preserving linguistic characteristics. This includes maintaining sentiments or specific keywords.\n- **Performance Metrics Analysis**: For performance-related KPI analysis, include methods for detecting outliers and normalizing scores to facilitate comparisons across different datasets or campaigns.\n- **Word Cloud Creation**: When generating word clouds, ensure to create distinct visual representations for different categories (questions vs. answers) and apply suitable color schemes to enhance differentiation.\n\n### Performance and Clarity:\n- Clean and preprocess data according to the details provided in the input descriptions.\n- Aim to visualize insights simply and clearly, emphasizing ease of understanding.\n- Strictly adhere to any specific instructions from the styling index, keeping the target audience's comprehension in mind when designing visual representations."}, - { - "template_name": "planner_sk_learn_agent", - "display_name": "Machine Learning Agent", - "description": "Multi-agent planner variant: Trains and evaluates machine learning models using scikit-learn, including classification, regression, and clustering with feature importance insights", - "icon_url": "/icons/templates/sk_learn_agent.svg", - "category": "Data Modelling", - "is_premium_only": false, - "variant_type": "planner", - "base_agent": "sk_learn_agent", - "is_active": true, - "prompt_template": "**Agent Definition:**\nYou are a machine learning agent in a multi-agent data analytics pipeline.\nYou are given:\n* A dataset (often cleaned and feature-engineered).\n* A user-defined goal (e.g., classification, regression, clustering).\n* Agent-specific **plan instructions** specifying:\n* Which **variables** you are expected to **CREATE** (e.g., `trained_model`, `predictions`).\n* Which **variables** you will **USE** (e.g., `df_cleaned`, `target_variable`, `feature_columns`).\n* A set of **instructions** outlining additional processing or handling for these variables (e.g., handling missing values, applying transformations, or other task-specific guidelines).\n**Your Responsibilities:**\n* Use the scikit-learn library to implement the appropriate ML pipeline.\n* Always split data into training and testing sets where applicable.\n* Use `print()` for all outputs.\n* Ensure your code is:\n* **Reproducible**: Set `random_state=42` wherever applicable.\n* **Modular**: Avoid deeply nested code.\n* **Focused on model building**, not visualization (leave plotting to the `data_viz_agent`).\n**You must not:**\n* Visualize anything (that's another agent's job).\n* Rely on hardcoded column names — use those passed via `plan_instructions`.\n* **Never create or modify any variables not explicitly mentioned in `plan_instructions['CREATE']`.**\n* **Never create the `df` variable**. You will **only** work with the variables passed via the `plan_instructions`.\n* Do not introduce intermediate variables unless they are listed in `plan_instructions['CREATE']`.\n**Instructions to Follow:**\n1. **CREATE** only the variables specified in the `plan_instructions['CREATE']` list. Do not create any intermediate or new variables.\n2. **USE** only the variables specified in the `plan_instructions['USE']` list. You are **not allowed** to create or modify any variables not listed in the plan instructions.\n3. Follow any **processing instructions** in the `plan_instructions['INSTRUCTIONS']` list. This might include tasks like handling missing values, scaling features, or encoding categorical variables. Always perform these steps on the variables specified in the `plan_instructions`.\n4. Do **not reassign or modify** any variables passed via `plan_instructions`. These should be used as-is.\n**Output:**\n* The **code** implementing the ML task, including all required steps.\n* A **summary** of what the model does, how it is evaluated, and why it fits the goal.\n* Respond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "feature_engineering_agent", - "display_name": "Feature Engineering Agent", - "description": "Advanced feature creation and selection for machine learning pipelines using various encoding and transformation techniques", - "icon_url": "/icons/templates/feature-engineering.png", - "category": "Data Modelling", - "is_premium_only": true, - "variant_type": "individual", - "base_agent": "feature_engineering_agent", - "is_active": true, - "prompt_template": "You are a feature engineering expert for machine learning pipelines. Your task is to take a dataset and a user-defined goal and create meaningful features that improve model performance.\n\nIMPORTANT Instructions:\n- Create meaningful features from raw data based on the user's goal\n- Apply feature scaling, encoding, and transformation techniques\n- Handle categorical variables with appropriate encoding methods (one-hot, label, target encoding)\n- Create polynomial features, interactions, and domain-specific features when beneficial\n- Perform feature selection using statistical and ML methods\n- Handle time-series feature engineering when applicable (lag features, rolling statistics)\n- Ensure features are robust and avoid data leakage\n- Use libraries like pandas, numpy, scikit-learn for feature engineering\n- Document feature engineering decisions and rationale\n\nProvide a concise bullet-point summary of the feature engineering operations performed.\n\nExample Summary:\n• Created 15 new features including polynomial interactions between price and quantity\n• Applied target encoding to categorical variables with high cardinality\n• Generated time-based features: day of week, month, rolling 7-day averages\n• Removed 8 highly correlated features (correlation > 0.95)\n• Applied StandardScaler to numerical features for model compatibility\n• Final feature set: 23 features with improved signal-to-noise ratio\n\nRespond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "planner_feature_engineering_agent", - "display_name": "Feature Engineering Agent", - "description": "Multi-agent planner variant: Advanced feature creation and selection for machine learning pipelines using various encoding and transformation techniques", - "icon_url": "/icons/templates/feature-engineering.png", - "category": "Data Modelling", - "is_premium_only": true, - "variant_type": "planner", - "base_agent": "feature_engineering_agent", - "is_active": true, - "prompt_template": "You are a feature engineering expert optimized for multi-agent data analytics pipelines.\n\nYou are given:\n* A dataset (often raw or lightly processed).\n* A user-defined goal (e.g., improve model performance, create specific feature types).\n* **plan_instructions** containing:\n * **'create'**: Variables you must create (e.g., ['engineered_features', 'feature_names', 'scaler_object'])\n * **'use'**: Variables you must use (e.g., ['raw_data', 'target_column'])\n * **'instruction'**: Specific feature engineering instructions\n\n### Your Planner-Optimized Responsibilities:\n* **ALWAYS follow plan_instructions** - essential for pipeline coordination\n* Create ONLY the variables specified in plan_instructions['create']\n* Use ONLY the variables specified in plan_instructions['use']\n* Apply feature engineering techniques as per plan_instructions['instruction']\n* Ensure engineered features integrate seamlessly with downstream ML agents\n\n### Feature Engineering Techniques:\n* Categorical encoding (one-hot, label, target encoding)\n* Numerical transformations (scaling, normalization, polynomial features)\n* Time-series features (lag features, rolling statistics, temporal patterns)\n* Feature selection and dimensionality reduction\n* Interaction features and domain-specific feature creation\n* Handle missing values and outliers appropriately\n\n### Multi-Agent Best Practices:\n* Use exact variable names from plan_instructions['create']\n* Ensure feature compatibility for downstream agents\n* Maintain data integrity and prevent leakage\n* Document feature engineering decisions for pipeline transparency\n\n### Output:\n* Python code implementing feature engineering per plan_instructions\n* Summary of features created and transformations applied\n* Focus on seamless integration with ML modeling agents\n\nRespond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "polars_agent", - "display_name": "Polars Agent", - "description": "High-performance data processing using Polars for large datasets with lazy evaluation and efficient memory usage", - "icon_url": "/icons/templates/polars_github_logo_rect_dark_name.svg", - "category": "Data Manipulation", - "is_premium_only": true, - "variant_type": "individual", - "base_agent": "polars_agent", - "is_active": true, - "prompt_template": "You are a Polars expert for high-performance data processing. Your task is to take a dataset and a user-defined goal and use Polars library for efficient data manipulation based on the user's goal.\n\nIMPORTANT Instructions:\n- Use Polars for efficient data manipulation and analysis\n- Leverage lazy evaluation for optimal performance with .lazy() and .collect()\n- Handle large datasets that don't fit in memory using streaming\n- Use Polars expressions (pl.col, pl.when, etc.) for complex transformations\n- Optimize query plans for speed and memory efficiency\n- Convert to/from pandas when needed for compatibility with other tools\n- Use appropriate data types to minimize memory usage\n- Apply Polars-specific optimizations like predicate pushdown\n- Focus on performance and memory efficiency over simplicity\n\nProvide a concise bullet-point summary of the Polars operations performed.\n\nExample Summary:\n• Processed 10M row dataset using lazy evaluation for memory efficiency\n• Applied complex filtering and aggregations with 5x speedup vs pandas\n• Used Polars expressions for vectorized string operations\n• Implemented window functions for time-series calculations\n• Optimized memory usage by selecting appropriate dtypes (reduced from 2GB to 500MB)\n• Final output: clean, aggregated dataset ready for analysis\n\nRespond in the user's language for all summary and reasoning but keep the code in english" - }, - { - "template_name": "planner_polars_agent", - "display_name": "Polars Agent", - "description": "Multi-agent planner variant: High-performance data processing using Polars for large datasets with lazy evaluation and efficient memory usage", - "icon_url": "https://raw.githubusercontent.com/pola-rs/polars-static/master/logos/polars_github_logo_rect_dark_name.svg", - "category": "Data Manipulation", - "is_premium_only": true, - "variant_type": "planner", - "base_agent": "polars_agent", - "is_active": true, - "prompt_template": "You are a Polars expert optimized for multi-agent data processing pipelines.\n\nYou are given:\n* A dataset (often large or complex).\n* A user-defined goal (e.g., data transformation, aggregation, filtering).\n* **plan_instructions** containing:\n * **'create'**: Variables you must create (e.g., ['processed_data', 'summary_stats'])\n * **'use'**: Variables you must use (e.g., ['raw_data', 'filter_conditions'])\n * **'instruction'**: Specific data processing instructions\n\n### Your Planner-Optimized Responsibilities:\n* **ALWAYS follow plan_instructions** - critical for pipeline data flow\n* Create ONLY the variables specified in plan_instructions['create']\n* Use ONLY the variables specified in plan_instructions['use']\n* Apply Polars operations as per plan_instructions['instruction']\n* Ensure processed data integrates seamlessly with downstream agents\n\n### Polars Optimization Techniques:\n* Use lazy evaluation (.lazy().collect()) for memory efficiency\n* Apply predicate pushdown and projection pushdown optimizations\n* Leverage Polars expressions for vectorized operations\n* Use appropriate data types to minimize memory footprint\n* Implement streaming for datasets larger than memory\n* Convert to pandas DataFrame only when required by downstream agents\n\n### Multi-Agent Best Practices:\n* Use exact variable names from plan_instructions['create']\n* Ensure data format compatibility for subsequent agents\n* Maintain data integrity and schema consistency\n* Optimize for both speed and memory usage in pipeline context\n\n### Output:\n* Python code implementing Polars operations per plan_instructions\n* Summary of data processing and optimizations applied\n* Focus on high-performance data flow in multi-agent pipeline\n\nRespond in the user's language for all summary and reasoning but keep the code in english" - } - ], - "remove": [] -} \ No newline at end of file diff --git a/app.py b/app.py index 03b022edb0cf3172e967c99fe32c104965f4ccf1..6733d0369b0e3cbd9294ab55d3ebba7b14dd5434 100644 --- a/app.py +++ b/app.py @@ -1,2888 +1,820 @@ # Standard library imports - import asyncio - import json - import logging - import os - import time - import uuid - from io import StringIO - from typing import List, Optional -import ast - -import markdown - -from bs4 import BeautifulSoup - -import pandas as pd - -from datetime import datetime, UTC - # Third-party imports - +import groq +import pandas as pd import uvicorn - from dotenv import load_dotenv - from fastapi import ( - Depends, - FastAPI, - File, - Form, - HTTPException, - Request, - UploadFile - ) - from fastapi.middleware.cors import CORSMiddleware - from fastapi.responses import JSONResponse, StreamingResponse - from fastapi.security import APIKeyHeader - +from llama_index.core import Document, VectorStoreIndex from pydantic import BaseModel - - # Local application imports - from scripts.format_response import format_response_to_markdown - from src.agents.agents import * - from src.agents.retrievers.retrievers import * - from src.managers.ai_manager import AI_Manager - from src.managers.session_manager import SessionManager - -from src.managers.app_manager import AppState - from src.routes.analytics_routes import router as analytics_router - -from src.routes.blog_routes import router as blog_router - from src.routes.chat_routes import router as chat_router - from src.routes.code_routes import router as code_router - -from src.routes.feedback_routes import router as feedback_router - from src.routes.session_routes import router as session_router, get_session_id_dependency - -from src.routes.deep_analysis_routes import router as deep_analysis_router - -from src.routes.templates_routes import router as templates_router - -from src.schemas.query_schema import QueryRequest - +from src.schemas.query_schemas import QueryRequest from src.utils.logger import Logger -from src.routes.session_routes import apply_model_safeguards - - -# Import deep analysis components directly - -# from src.agents.try_deep_agents import deep_analysis_module - -from src.agents.deep_agents import deep_analysis_module - -from src.utils.generate_report import generate_html_report - - - -from src.utils.model_registry import MODEL_OBJECTS - - - -logger = Logger("app", see_time=True, console_log=True) +logger = Logger("app", see_time=True, console_log=False) load_dotenv() - - -# Request models - -class DeepAnalysisRequest(BaseModel): - - goal: str - - - -class DeepAnalysisResponse(BaseModel): - - goal: str - - deep_questions: str - - deep_plan: str - - summaries: List[str] - - code: str - - plotly_figs: List - - synthesis: List[str] - - final_conclusion: str - - html_report: Optional[str] = None - - - -styling_instructions = [ - { - "category": "line_charts", - "description": "Used to visualize trends and changes over time, often with multiple series.", - "styling": { - "template": "plotly_white", - "axes_line_width": 0.2, - "grid_width": 1, - "title": { - "bold_html": True, - "include": True - }, - "colors": "use multiple colors if more than one line", - "annotations": ["min", "max"], - "number_format": { - "apply_k_m": True, - "thresholds": {"K": 1000, "M": 100000}, - "percentage_decimals": 2, - "percentage_sign": True - }, - "default_size": {"height": 1200, "width": 1000} - } - }, - { - "category": "bar_charts", - "description": "Useful for comparing discrete categories or groups with bars representing values.", - "styling": { - "template": "plotly_white", - "axes_line_width": 0.2, - "grid_width": 1, - "title": {"bold_html": True, "include": True}, - "annotations": ["bar values"], - "number_format": { - "apply_k_m": True, - "thresholds": {"K": 1000, "M": 100000}, - "percentage_decimals": 2, - "percentage_sign": True - }, - "default_size": {"height": 1200, "width": 1000} - } - }, - { - "category": "histograms", - "description": "Display the distribution of a data set, useful for returns or frequency distributions.", - "styling": { - "template": "plotly_white", - "bin_size": 50, - "axes_line_width": 0.2, - "grid_width": 1, - "title": {"bold_html": True, "include": True}, - "annotations": ["x values"], - "number_format": { - "apply_k_m": True, - "thresholds": {"K": 1000, "M": 100000}, - "percentage_decimals": 2, - "percentage_sign": True - }, - "default_size": {"height": 1200, "width": 1000} - } - }, - { - "category": "pie_charts", - "description": "Show composition or parts of a whole with slices representing categories.", - "styling": { - "template": "plotly_white", - "top_categories_to_show": 10, - "bundle_rest_as": "Others", - "axes_line_width": 0.2, - "grid_width": 1, - "title": {"bold_html": True, "include": True}, - "annotations": ["x values"], - "number_format": { - "apply_k_m": True, - "thresholds": {"K": 1000, "M": 100000}, - "percentage_decimals": 2, - "percentage_sign": True - }, - "default_size": {"height": 1200, "width": 1000} - } - }, - { - "category": "tabular_and_generic_charts", - "description": "Applies to charts where number formatting needs flexibility, including mixed or raw data.", - "styling": { - "template": "plotly_white", - "axes_line_width": 0.2, - "grid_width": 1, - "title": {"bold_html": True, "include": True}, - "annotations": ["x values"], - "number_format": { - "apply_k_m": True, - "thresholds": {"K": 1000, "M": 100000}, - "exclude_if_commas_present": True, - "exclude_if_not_numeric": True, - "percentage_decimals": 2, - "percentage_sign": True - }, - "default_size": {"height": 1200, "width": 1000} - } - }, - { - "category": "heat_maps", - "description": "Show data density or intensity using color scales on a matrix or grid.", - "styling": { - "template": "plotly_white", - "axes_styles": { - "line_color": "black", - "line_width": 0.2, - "grid_width": 1, - "format_numbers_as_k_m": True, - "exclude_non_numeric_formatting": True - }, - "title": {"bold_html": True, "include": True}, - "default_size": {"height": 1200, "width": 1000} - } - }, - { - "category": "histogram_distribution", - "description": "Specialized histogram for return distributions with opacity control.", - "styling": { - "template": "plotly_white", - "opacity": 0.75, - "axes_styles": { - "grid_width": 1, - "format_numbers_as_k_m": True, - "exclude_non_numeric_formatting": True - }, - "title": {"bold_html": True, "include": True}, - "default_size": {"height": 1200, "width": 1000} - } - } +styling_instructions = [ + """ + Dont ignore any of these instructions. + For a line chart always use plotly_white template, reduce x axes & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label and try to use multiple colors if more than one line + Annotate the min and max of the line + Display numbers in thousand(K) or Million(M) if larger than 1000/100000 + Show percentages in 2 decimal points with '%' sign + Default size of chart should be height =1200 and width =1000 + + """ + + , """ + Dont ignore any of these instructions. + For a bar chart always use plotly_white template, reduce x axes & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label + Always display numbers in thousand(K) or Million(M) if larger than 1000/100000. + Annotate the values of the bar chart + If variable is a percentage show in 2 decimal points with '%' sign. + Default size of chart should be height =1200 and width =1000 + """ + , + + """ + For a histogram chart choose a bin_size of 50 + Do not ignore any of these instructions + always use plotly_white template, reduce x & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label + Always display numbers in thousand(K) or Million(M) if larger than 1000/100000. Add annotations x values + If variable is a percentage show in 2 decimal points with '%' + Default size of chart should be height =1200 and width =1000 + """, + + + """ + For a pie chart only show top 10 categories, bundle rest as others + Do not ignore any of these instructions + always use plotly_white template, reduce x & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label + Always display numbers in thousand(K) or Million(M) if larger than 1000/100000. Add annotations x values + If variable is a percentage show in 2 decimal points with '%' + Default size of chart should be height =1200 and width =1000 + """, + + """ + Do not ignore any of these instructions + always use plotly_white template, reduce x & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label + Always display numbers in thousand(K) or Million(M) if larger than 1000/100000. Add annotations x values + Don't add K/M if number already in , or value is not a number + If variable is a percentage show in 2 decimal points with '%' + Default size of chart should be height =1200 and width =1000 + """, +""" + For a heat map + Use the 'plotly_white' template for a clean, white background. + Set a chart title + Style the X-axis with a black line color, 0.2 line width, 1 grid width, format 1000/1000000 as K/M + Do not format non-numerical numbers + .style the Y-axis with a black line color, 0.2 line width, 1 grid width format 1000/1000000 as K/M + Do not format non-numerical numbers + + . Set the figure dimensions to a height of 1200 pixels and a width of 1000 pixels. +""", +""" + For a Histogram, used for returns/distribution plotting + Use the 'plotly_white' template for a clean, white background. + Set a chart title + Style the X-axis 1 grid width, format 1000/1000000 as K/M + Do not format non-numerical numbers + .style the Y-axis, 1 grid width format 1000/1000000 as K/M + Do not format non-numerical numbers + + Use an opacity of 0.75 + + Set the figure dimensions to a height of 1200 pixels and a width of 1000 pixels. +""" ] -# Convert to list of JSON strings -styling_instructions = [str(chart_dict) for chart_dict in styling_instructions] - -# Output (just show first 2 for readability) - - - - # Add near the top of the file, after imports - DEFAULT_MODEL_CONFIG = { "provider": os.getenv("MODEL_PROVIDER", "openai"), - "model": os.getenv("MODEL_NAME", "gpt-5-mini"), + "model": os.getenv("MODEL_NAME", "gpt-4o-mini"), "api_key": os.getenv("OPENAI_API_KEY"), - "temperature": min(1.0, max(0.0, float(os.getenv("TEMPERATURE", "1.0")))), # Clamp to 0..1 - "max_tokens": int(os.getenv("MAX_TOKENS", 6000)), "cache": False - + "temperature": float(os.getenv("TEMPERATURE", 1.0)), + "max_tokens": int(os.getenv("MAX_TOKENS", 6000)) } - - # Create default LM config but don't set it globally - - - -default_lm = MODEL_OBJECTS[DEFAULT_MODEL_CONFIG['model']] - - - - - - - -dspy.configure(lm=default_lm, async_max_workers=1000) - - +if DEFAULT_MODEL_CONFIG["provider"].lower() == "groq": + default_lm = dspy.GROQ( + model=DEFAULT_MODEL_CONFIG["model"], + api_key=DEFAULT_MODEL_CONFIG["api_key"], + temperature=DEFAULT_MODEL_CONFIG["temperature"], + max_tokens=DEFAULT_MODEL_CONFIG["max_tokens"] + ) +elif DEFAULT_MODEL_CONFIG["provider"].lower() == "gemini": + default_lm = dspy.LM( + model=f"gemini/{DEFAULT_MODEL_CONFIG['model']}", + api_key=DEFAULT_MODEL_CONFIG["api_key"], + temperature=DEFAULT_MODEL_CONFIG["temperature"], + max_tokens=DEFAULT_MODEL_CONFIG["max_tokens"] + ) +else: + default_lm = dspy.LM( + model=DEFAULT_MODEL_CONFIG["model"], + api_key=DEFAULT_MODEL_CONFIG["api_key"], + temperature=DEFAULT_MODEL_CONFIG["temperature"], + max_tokens=DEFAULT_MODEL_CONFIG["max_tokens"] + ) # Function to get model config from session or use default - def get_session_lm(session_state): - """Get the appropriate LM instance for a session, or default if not configured""" - # First check if we have a valid session-specific model config - if session_state and isinstance(session_state, dict) and "model_config" in session_state: - model_config = session_state["model_config"] - if model_config and isinstance(model_config, dict) and "model" in model_config: - # Found valid session-specific model config, use it - provider = model_config.get("provider", "openai").lower() - - model_name = model_config.get("model", DEFAULT_MODEL_CONFIG["model"]) - - # Import and apply centralized safeguards (temperature + max_tokens) - - - requested_temp = model_config.get("temperature", DEFAULT_MODEL_CONFIG["temperature"]) - requested_max = model_config.get("max_tokens", DEFAULT_MODEL_CONFIG["max_tokens"]) - - safe_params = apply_model_safeguards( - model_name=model_name, - provider=provider, - temperature=requested_temp, - max_tokens=requested_max, - ) - - # Apply the safeguarded parameters - MODEL_OBJECTS[model_name].__dict__['kwargs']['max_tokens'] = safe_params["max_tokens"] - MODEL_OBJECTS[model_name].__dict__['kwargs']['temperature'] = safe_params["temperature"] - - - + if provider == "groq": + return dspy.GROQ( + model=model_config.get("model", DEFAULT_MODEL_CONFIG["model"]), + api_key=model_config.get("api_key", DEFAULT_MODEL_CONFIG["api_key"]), + temperature=model_config.get("temperature", DEFAULT_MODEL_CONFIG["temperature"]), + max_tokens=model_config.get("max_tokens", DEFAULT_MODEL_CONFIG["max_tokens"]) + ) + elif provider == "anthropic": + return dspy.LM( + model=model_config.get("model", DEFAULT_MODEL_CONFIG["model"]), + api_key=model_config.get("api_key", DEFAULT_MODEL_CONFIG["api_key"]), + temperature=model_config.get("temperature", DEFAULT_MODEL_CONFIG["temperature"]), + max_tokens=model_config.get("max_tokens", DEFAULT_MODEL_CONFIG["max_tokens"]) + ) + elif provider == "gemini": + return dspy.LM( + model=f"gemini/{model_config.get('model', DEFAULT_MODEL_CONFIG['model'])}", + api_key=model_config.get("api_key", DEFAULT_MODEL_CONFIG["api_key"]), + temperature=model_config.get("temperature", DEFAULT_MODEL_CONFIG["temperature"]), + max_tokens=model_config.get("max_tokens", DEFAULT_MODEL_CONFIG["max_tokens"]) + ) + else: # OpenAI is the default + return dspy.LM( + model=model_config.get("model", DEFAULT_MODEL_CONFIG["model"]), + api_key=model_config.get("api_key", DEFAULT_MODEL_CONFIG["api_key"]), + temperature=model_config.get("temperature", DEFAULT_MODEL_CONFIG["temperature"]), + max_tokens=model_config.get("max_tokens", DEFAULT_MODEL_CONFIG["max_tokens"]) + ) - # If no valid session config, use default - - return MODEL_OBJECTS[model_name] - - + return default_lm # Initialize retrievers with empty data first - - +def initialize_retrievers(styling_instructions: List[str], doc: List[str]): + try: + style_index = VectorStoreIndex.from_documents([Document(text=x) for x in styling_instructions]) + data_index = VectorStoreIndex.from_documents([Document(text=x) for x in doc]) + return {"style_index": style_index, "dataframe_index": data_index} + except Exception as e: + logger.log_message(f"Error initializing retrievers: {str(e)}", level=logging.ERROR) + raise e # clear console - def clear_console(): - os.system('cls' if os.name == 'nt' else 'clear') - - - # Check for Housing.csv - housing_csv_path = "Housing.csv" - if not os.path.exists(housing_csv_path): - logger.log_message(f"Housing.csv not found at {os.path.abspath(housing_csv_path)}", level=logging.ERROR) - raise FileNotFoundError(f"Housing.csv not found at {os.path.abspath(housing_csv_path)}") +AVAILABLE_AGENTS = { + "data_viz_agent": data_viz_agent, + "sk_learn_agent": sk_learn_agent, + "statistical_analytics_agent": statistical_analytics_agent, + "preprocessing_agent": preprocessing_agent, +} - -# All agents are now loaded from database - no hardcoded dictionaries needed - - +PLANNER_AGENTS = { + "planner_preprocessing_agent": planner_preprocessing_agent, + "planner_sk_learn_agent": planner_sk_learn_agent, + "planner_statistical_analytics_agent": planner_statistical_analytics_agent, + "planner_data_viz_agent": planner_data_viz_agent, +} # Add session header - X_SESSION_ID = APIKeyHeader(name="X-Session-ID", auto_error=False) - - # Update AppState class to use SessionManager - -# The AppState class is now in src.managers.app_manager - - - -# Initialize FastAPI app with state - -app = FastAPI(title="AI Analytics API", version="1.0") - -# Pass required parameters to AppState -app.state = AppState(styling_instructions, chat_history_name_agent, DEFAULT_MODEL_CONFIG) - - - - - -# Configure middleware - -# Use a wildcard for local development or read from environment - -is_development = os.getenv("ENVIRONMENT", "development").lower() == "development" - - - -allowed_origins = [] - -frontend_url = os.getenv("FRONTEND_URL", "").strip() - -print(f"FRONTEND_URL: {frontend_url}") - -if is_development: - - allowed_origins = ["*"] - -elif frontend_url: - - allowed_origins = [frontend_url] - -else: - - logger.log_message("CORS misconfigured: FRONTEND_URL not set", level=logging.ERROR) - - allowed_origins = [] # or set a default safe origin - - - -# Add a strict origin verification middleware - -@app.middleware("http") - -async def verify_origin_middleware(request: Request, call_next): - - # Skip origin check in development mode - - if is_development: - - return await call_next(request) - +class AppState: + def __init__(self): + self._session_manager = SessionManager(styling_instructions, PLANNER_AGENTS) + self.model_config = DEFAULT_MODEL_CONFIG.copy() + # Update the SessionManager with the current model_config + self._session_manager._app_model_config = self.model_config + self.ai_manager = AI_Manager() + self.chat_name_agent = chat_history_name_agent + def get_session_state(self, session_id: str): + """Get or create session-specific state using the SessionManager""" + return self._session_manager.get_session_state(session_id) - # Get the origin from the request headers + def clear_session_state(self, session_id: str): + """Clear session-specific state using the SessionManager""" + self._session_manager.clear_session_state(session_id) - origin = request.headers.get("origin") + def update_session_dataset(self, session_id: str, df, name, desc): + """Update dataset for a specific session using the SessionManager""" + self._session_manager.update_session_dataset(session_id, df, name, desc) + def reset_session_to_default(self, session_id: str): + """Reset a session to use the default dataset using the SessionManager""" + self._session_manager.reset_session_to_default(session_id) - - # Log the origin for debugging - - if origin: - - print(f"Request from origin: {origin}") - + def set_session_user(self, session_id: str, user_id: int, chat_id: int = None): + """Associate a user with a session using the SessionManager""" + return self._session_manager.set_session_user(session_id, user_id, chat_id) - - # If no origin header or origin not in allowed list, reject the request - - if origin and frontend_url and origin != frontend_url: - - print(f"Blocked request from unauthorized origin: {origin}") - - return JSONResponse( - - status_code=403, - - content={"detail": "Not authorized"} - - ) - + def get_ai_manager(self): + """Get the AI Manager instance""" + return self.ai_manager + def get_provider_for_model(self, model_name): + return self.ai_manager.get_provider_for_model(model_name) + + def calculate_cost(self, model_name, input_tokens, output_tokens): + return self.ai_manager.calculate_cost(model_name, input_tokens, output_tokens) + + def save_usage_to_db(self, user_id, chat_id, model_name, provider, prompt_tokens, completion_tokens, total_tokens, query_size, response_size, cost, request_time_ms, is_streaming=False): + return self.ai_manager.save_usage_to_db(user_id, chat_id, model_name, provider, prompt_tokens, completion_tokens, total_tokens, query_size, response_size, round(cost, 7), request_time_ms, is_streaming) + + def get_tokenizer(self): + return self.ai_manager.tokenizer + + def get_chat_history_name_agent(self): + return dspy.Predict(self.chat_name_agent) - # Continue processing the request if origin is allowed - - return await call_next(request) - - - -# CORS middleware (still needed for browser preflight) +# Initialize FastAPI app with state +app = FastAPI(title="AI Analytics API", version="1.0") +app.state = AppState() +# Configure middleware app.add_middleware( - CORSMiddleware, - - allow_origins=allowed_origins, - - allow_origin_regex=None, - + allow_origins=["*"], allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], - - expose_headers=["*"], - - max_age=600 # Cache preflight requests for 10 minutes (for performance) - + expose_headers=["Content-Type", "Content-Length"] ) - - # Add these constants at the top of the file with other imports/constants - RESPONSE_ERROR_INVALID_QUERY = "Please provide a valid query..." - RESPONSE_ERROR_NO_DATASET = "No dataset is currently loaded. Please link a dataset before proceeding with your analysis." - DEFAULT_TOKEN_RATIO = 1.5 - -REQUEST_TIMEOUT_SECONDS = 90 # Timeout for LLM requests - -MAX_RECENT_MESSAGES = 5 - +REQUEST_TIMEOUT_SECONDS = 60 # Timeout for LLM requests +MAX_RECENT_MESSAGES = 3 DB_BATCH_SIZE = 10 # For future batch DB operations - - +# Replace the existing chat_with_agent function @app.post("/chat/{agent_name}", response_model=dict) - async def chat_with_agent( - agent_name: str, - request: QueryRequest, - request_obj: Request, - session_id: str = Depends(get_session_id_dependency) - ): - session_state = app.state.get_session_state(session_id) - - logger.log_message(f"[DEBUG] chat_with_agent called with agent: '{agent_name}', query: '{request.query[:100]}...'", level=logging.DEBUG) - - try: - # Extract and validate query parameters - - logger.log_message(f"[DEBUG] Updating session from query params", level=logging.DEBUG) - _update_session_from_query_params(request_obj, session_state) - - logger.log_message(f"[DEBUG] Session state after query params: user_id={session_state.get('user_id')}, chat_id={session_state.get('chat_id')}", level=logging.DEBUG) - - # Validate dataset and agent name + if session_state["current_df"] is None: + raise HTTPException(status_code=400, detail=RESPONSE_ERROR_NO_DATASET) - if session_state["datasets"] is None: - logger.log_message(f"[DEBUG] No dataset loaded", level=logging.DEBUG) - - raise HTTPException(status_code=400, detail=RESPONSE_ERROR_NO_DATASET) - - - - # Log the dataset being used for analysis with detailed information - datasets = session_state["datasets"] - dataset_names = list(datasets.keys()) - if dataset_names: - current_dataset_name = dataset_names[-1] # Get the last (most recent) dataset - dataset_shape = datasets[current_dataset_name].shape - - # Check if this is the default dataset and explain why - session_name = session_state.get("name", "") - is_default_dataset = (current_dataset_name == "df" and session_name == "Housing.csv") or current_dataset_name == "Housing.csv" - - if is_default_dataset: - logger.log_message(f"[ANALYSIS] Using DEFAULT dataset 'Housing.csv' for analysis (shape: {dataset_shape[0]} rows, {dataset_shape[1]} columns)", level=logging.INFO) - logger.log_message(f"[ANALYSIS] Reason: No custom dataset uploaded yet - using default Housing.csv dataset", level=logging.INFO) - else: - logger.log_message(f"[ANALYSIS] Using CUSTOM dataset '{current_dataset_name}' for analysis (shape: {dataset_shape[0]} rows, {dataset_shape[1]} columns)", level=logging.INFO) - logger.log_message(f"[ANALYSIS] This is a user-uploaded dataset, not the default", level=logging.INFO) - else: - logger.log_message(f"[ANALYSIS] No datasets available in session {session_id}", level=logging.WARNING) - - logger.log_message(f"[DEBUG] About to validate agent name: '{agent_name}'", level=logging.DEBUG) - - _validate_agent_name(agent_name, session_state) - - logger.log_message(f"[DEBUG] Agent validation completed successfully", level=logging.DEBUG) - + _validate_agent_name(agent_name) - # Record start time for timing - start_time = time.time() - - # Get chat context and prepare query - - logger.log_message(f"[DEBUG] Preparing query with context", level=logging.DEBUG) - enhanced_query = _prepare_query_with_context(request.query, session_state) - - logger.log_message(f"[DEBUG] Enhanced query length: {len(enhanced_query)}", level=logging.DEBUG) - - - # Initialize agent - handle standard, template, and custom agents - + # Initialize agent if "," in agent_name: - - logger.log_message(f"[DEBUG] Processing multiple agents: {agent_name}", level=logging.DEBUG) - - # Multiple agents case - - agent_list = [agent.strip() for agent in agent_name.split(",")] - - - - # Categorize agents - - standard_agents = [agent for agent in agent_list if _is_standard_agent(agent)] - - template_agents = [agent for agent in agent_list if _is_template_agent(agent)] - - custom_agents = [agent for agent in agent_list if not _is_standard_agent(agent) and not _is_template_agent(agent)] - - - - logger.log_message(f"[DEBUG] Agent categorization - standard: {standard_agents}, template: {template_agents}, custom: {custom_agents}", level=logging.DEBUG) - - - - if custom_agents: - - # If any custom agents, use session AI system for all - - ai_system = session_state["ai_system"] - - session_lm = get_session_lm(session_state) - - logger.log_message(f"[DEBUG] Using custom agent execution path", level=logging.DEBUG) - - with dspy.context(lm=session_lm): - - response = await asyncio.wait_for( - - _execute_custom_agents(ai_system, agent_list, enhanced_query), - - timeout=REQUEST_TIMEOUT_SECONDS - - ) - - logger.log_message(f"[DEBUG] Custom agents response type: {type(response)}, keys: {list(response.keys()) if isinstance(response, dict) else 'not a dict'}", level=logging.DEBUG) - - else: - - # All standard/template agents - use auto_analyst_ind which loads from DB - - user_id = session_state.get("user_id") - - logger.log_message(f"[DEBUG] Using auto_analyst_ind for multiple standard/template agents with user_id: {user_id}", level=logging.DEBUG) - - - - # Create database session for agent loading - - from src.db.init_db import session_factory - - db_session = session_factory() - - try: - - # auto_analyst_ind will load all agents from database - - logger.log_message(f"[DEBUG] Creating auto_analyst_ind instance", level=logging.DEBUG) - - agent = auto_analyst_ind(agents=[], retrievers=session_state["retrievers"], user_id=user_id, db_session=db_session) - - session_lm = get_session_lm(session_state) - - logger.log_message(f"[DEBUG] About to call agent.forward with query and agent list", level=logging.DEBUG) - - with dspy.context(lm=session_lm): - - response = await asyncio.wait_for( - - agent(enhanced_query, ",".join(agent_list)), - - timeout=REQUEST_TIMEOUT_SECONDS - - ) - - logger.log_message(f"[DEBUG] auto_analyst_ind response type: {type(response)}, content: {str(response)[:200]}...", level=logging.DEBUG) - - finally: - - db_session.close() - + agent_list = [AVAILABLE_AGENTS[agent.strip()] for agent in agent_name.split(",")] + agent = auto_analyst_ind(agents=agent_list, retrievers=session_state["retrievers"]) else: - - logger.log_message(f"[DEBUG] Processing single agent: {agent_name}", level=logging.DEBUG) - - # Single agent case - - if _is_standard_agent(agent_name) or _is_template_agent(agent_name): - - # Standard or template agent - use auto_analyst_ind which loads from DB - - user_id = session_state.get("user_id") - - logger.log_message(f"[DEBUG] Using auto_analyst_ind for single standard/template agent '{agent_name}' with user_id: {user_id}", level=logging.DEBUG) - - - - # Create database session for agent loading - - from src.db.init_db import session_factory - - db_session = session_factory() - - try: - - # auto_analyst_ind will load all agents from database - - logger.log_message(f"[DEBUG] Creating auto_analyst_ind instance for single agent", level=logging.DEBUG) - - agent = auto_analyst_ind(agents=[], retrievers=session_state["retrievers"], user_id=user_id, db_session=db_session) - - session_lm = get_session_lm(session_state) - - logger.log_message(f"[DEBUG] About to call agent.forward for single agent '{agent_name}'", level=logging.DEBUG) - - with dspy.context(lm=session_lm): - - response = await asyncio.wait_for( - - agent(enhanced_query, agent_name), - - timeout=REQUEST_TIMEOUT_SECONDS - - ) - - logger.log_message(f"[DEBUG] Single agent response type: {type(response)}, content: {str(response)[:200]}...", level=logging.DEBUG) - - finally: - - db_session.close() - - else: - - # Custom agent - use session AI system - - ai_system = session_state["ai_system"] - - session_lm = get_session_lm(session_state) - - logger.log_message(f"[DEBUG] Using custom agent execution for '{agent_name}'", level=logging.DEBUG) - - with dspy.context(lm=session_lm): - - response = await asyncio.wait_for( - - _execute_custom_agents(ai_system, [agent_name], enhanced_query), - - timeout=REQUEST_TIMEOUT_SECONDS - - ) - - logger.log_message(f"[DEBUG] Custom single agent response type: {type(response)}, content: {str(response)[:200]}...", level=logging.DEBUG) - + agent = auto_analyst_ind(agents=[AVAILABLE_AGENTS[agent_name]], retrievers=session_state["retrievers"]) - - logger.log_message(f"[DEBUG] About to format response to markdown. Response type: {type(response)}", level=logging.DEBUG) - - formatted_response = format_response_to_markdown(response, agent_name, session_state["datasets"]) - logger.log_message(f"[DEBUG] Formatted response type: {type(formatted_response)}, length: {len(str(formatted_response))}", level=logging.DEBUG) - + # Execute agent with timeout + try: + # Get session-specific model + session_lm = get_session_lm(session_state) + + # Use session-specific model for this request + with dspy.context(lm=session_lm): + response = await asyncio.wait_for( + asyncio.to_thread(agent, enhanced_query, agent_name), + timeout=REQUEST_TIMEOUT_SECONDS + ) + except asyncio.TimeoutError: + logger.log_message(f"Agent execution timed out for {agent_name}", level=logging.WARNING) + raise HTTPException(status_code=504, detail="Request timed out. Please try a simpler query.") + except Exception as agent_error: + logger.log_message(f"Agent execution failed: {str(agent_error)}", level=logging.ERROR) + raise HTTPException(status_code=500, detail="Failed to process query. Please try again.") + + formatted_response = format_response_to_markdown(response, agent_name, session_state["current_df"]) - if formatted_response == RESPONSE_ERROR_INVALID_QUERY: - - logger.log_message(f"[DEBUG] Response was invalid query error", level=logging.DEBUG) - return { - "agent_name": agent_name, - "query": request.query, - "response": formatted_response, - "session_id": session_id - } - - # Track usage statistics - if session_state.get("user_id"): - - logger.log_message(f"[DEBUG] Tracking model usage", level=logging.DEBUG) - _track_model_usage( - session_state=session_state, - enhanced_query=enhanced_query, - response=response, - processing_time_ms=int((time.time() - start_time) * 1000) - ) - - - logger.log_message(f"[DEBUG] chat_with_agent completed successfully", level=logging.DEBUG) - return { - "agent_name": agent_name, - "query": request.query, # Return original query without context - "response": formatted_response, - "session_id": session_id - } - except HTTPException: - # Re-raise HTTP exceptions to preserve status codes - - logger.log_message(f"[DEBUG] HTTPException caught and re-raised", level=logging.DEBUG) - raise - - except asyncio.TimeoutError: - - logger.log_message(f"[ERROR] Timeout error in chat_with_agent", level=logging.ERROR) - - raise HTTPException(status_code=504, detail="Request timed out. Please try a simpler query.") - except Exception as e: - - logger.log_message(f"[ERROR] Unexpected error in chat_with_agent: {str(e)}", level=logging.ERROR) - - logger.log_message(f"[ERROR] Exception type: {type(e)}, traceback: {str(e)}", level=logging.ERROR) - - import traceback - - logger.log_message(f"[ERROR] Full traceback: {traceback.format_exc()}", level=logging.ERROR) - + logger.log_message(f"Unexpected error in chat_with_agent: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=500, detail="An unexpected error occurred. Please try again later.") - - - - - + + @app.post("/chat", response_model=dict) - async def chat_with_all( - request: QueryRequest, - request_obj: Request, - session_id: str = Depends(get_session_id_dependency) - ): - session_state = app.state.get_session_state(session_id) - - try: - # Extract and validate query parameters - _update_session_from_query_params(request_obj, session_state) - - # Validate dataset - - if session_state["datasets"] is None: + if session_state["current_df"] is None: raise HTTPException(status_code=400, detail=RESPONSE_ERROR_NO_DATASET) - - if session_state["ai_system"] is None: - raise HTTPException(status_code=500, detail="AI system not properly initialized.") - - # Get session-specific model - session_lm = get_session_lm(session_state) - - # Create streaming response - return StreamingResponse( - _generate_streaming_responses(session_state, request.query, session_lm), - media_type='text/event-stream', - headers={ - 'Cache-Control': 'no-cache', - 'Connection': 'keep-alive', - 'Content-Type': 'text/event-stream', - 'Access-Control-Allow-Origin': '*', - 'X-Accel-Buffering': 'no' - } - ) - except HTTPException: - # Re-raise HTTP exceptions to preserve status codes - raise - except Exception as e: - + logger.log_message(f"Unexpected error in chat_with_all: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=500, detail="An unexpected error occurred. Please try again later.") - - - # Helper functions to reduce duplication and improve modularity - def _update_session_from_query_params(request_obj: Request, session_state: dict): - """Extract and validate chat_id and user_id from query parameters""" - # Check for chat_id in query parameters - if "chat_id" in request_obj.query_params: - try: - chat_id_param = int(request_obj.query_params.get("chat_id")) - # Update session state with this chat ID - session_state["chat_id"] = chat_id_param - except (ValueError, TypeError): - logger.log_message("Invalid chat_id parameter", level=logging.WARNING) - # Continue without updating chat_id - - # Check for user_id in query parameters - if "user_id" in request_obj.query_params: - try: - user_id = int(request_obj.query_params["user_id"]) - session_state["user_id"] = user_id - except (ValueError, TypeError): - raise HTTPException( - status_code=400, - detail="Invalid user_id in query params. Please provide a valid integer." - ) - - - -def _validate_agent_name(agent_name: str, session_state: dict = None): - - """Validate that the agent name(s) are available""" - - logger.log_message(f"[DEBUG] Validating agent name: '{agent_name}'", level=logging.DEBUG) - - - +def _validate_agent_name(agent_name: str): + """Validate that the requested agent(s) exist""" if "," in agent_name: - - # Multiple agents - agent_list = [agent.strip() for agent in agent_name.split(",")] - - logger.log_message(f"[DEBUG] Multiple agents detected: {agent_list}", level=logging.DEBUG) - for agent in agent_list: - - is_available = _is_agent_available(agent, session_state) - - logger.log_message(f"[DEBUG] Agent '{agent}' availability: {is_available}", level=logging.DEBUG) - - if not is_available: - - available_agents = _get_available_agents_list(session_state) - - logger.log_message(f"[DEBUG] Agent '{agent}' not found. Available: {available_agents}", level=logging.DEBUG) - + if agent not in AVAILABLE_AGENTS: + available = list(AVAILABLE_AGENTS.keys()) raise HTTPException( - - status_code=400, - - detail=f"Agent '{agent}' not found. Available agents: {available_agents}" - + status_code=404, + detail=f"Agent '{agent}' not found. Available agents: {available}" ) - - else: - - # Single agent - - is_available = _is_agent_available(agent_name, session_state) - - logger.log_message(f"[DEBUG] Single agent '{agent_name}' availability: {is_available}", level=logging.DEBUG) - - if not is_available: - - available_agents = _get_available_agents_list(session_state) - - logger.log_message(f"[DEBUG] Agent '{agent_name}' not found. Available: {available_agents}", level=logging.DEBUG) - - raise HTTPException( - - status_code=400, - - detail=f"Agent '{agent_name}' not found. Available agents: {available_agents}" - - ) - - - - logger.log_message(f"[DEBUG] Agent validation passed for: '{agent_name}'", level=logging.DEBUG) - - - -def _is_agent_available(agent_name: str, session_state: dict = None) -> bool: - - """Check if an agent is available (standard, template, or custom)""" - - # Check if it's a standard agent - - if _is_standard_agent(agent_name): - - return True - - - - # Check if it's a template agent - - if _is_template_agent(agent_name): - - return True - - - - # Check if it's a custom agent in session - - if session_state and "ai_system" in session_state: - - ai_system = session_state["ai_system"] - - if hasattr(ai_system, 'agents') and agent_name in ai_system.agents: - - return True - - - - return False - - - -def _get_available_agents_list(session_state: dict = None) -> list: - - """Get list of all available agents from database""" - - from src.db.init_db import session_factory - - from src.agents.agents import load_all_available_templates_from_db - - - - # Core agents (always available) - - available = ["preprocessing_agent", "statistical_analytics_agent", "sk_learn_agent", "data_viz_agent"] - - - - # Add template agents from database - - db_session = session_factory() - - try: - - template_agents_dict = load_all_available_templates_from_db(db_session) - - # template_agents_dict is a dict with template_name as keys - - template_names = [template_name for template_name in template_agents_dict.keys() - - if template_name not in available and template_name != 'basic_qa_agent'] - - available.extend(template_names) - - except Exception as e: - - logger.log_message(f"Error loading template agents: {str(e)}", level=logging.ERROR) - - finally: - - db_session.close() - - - - return available - - - -def _is_standard_agent(agent_name: str) -> bool: - - """Check if agent is one of the 4 core standard agents""" - - standard_agents = ["preprocessing_agent", "statistical_analytics_agent", "sk_learn_agent", "data_viz_agent"] - - return agent_name in standard_agents - - - -def _is_template_agent(agent_name: str) -> bool: - - """Check if agent is a template agent""" - - try: - - from src.db.init_db import session_factory - - from src.db.schemas.models import AgentTemplate - - - - db_session = session_factory() - - try: - - template = db_session.query(AgentTemplate).filter( - - AgentTemplate.template_name == agent_name, - - AgentTemplate.is_active == True - - ).first() - - return template is not None - - finally: - - db_session.close() - - except Exception as e: - - logger.log_message(f"Error checking if {agent_name} is template: {str(e)}", level=logging.ERROR) - - return False - - - -async def _execute_custom_agents(ai_system, agent_names: list, query: str): - - """Execute custom agents using the session's AI system""" - - try: - - # For custom agents, we need to use the AI system's execute_agent method - - - - agent_results = [ai_system] - - if len(agent_names) == 1: - - # Single custom agent - - agent_name = agent_names[0] - - # Prepare inputs for the custom agent (similar to standard agents like data_viz_agent) - - dict_ = {} - - dict_['dataset'] = ai_system.dataset.retrieve(query)[0].text - - dict_['styling_index'] = ai_system.styling_index.retrieve(query)[0].text - - dict_['goal'] = query - - dict_['Agent_desc'] = str(ai_system.agent_desc) - - - - # Get input fields for this agent - - if agent_name in ai_system.agent_inputs: - - inputs = {x: dict_[x] for x in ai_system.agent_inputs[agent_name] if x in dict_} - - - - # Execute the custom agent - - agent_name_result, result_dict = await ai_system.agents[agent_name](**inputs) - - return {agent_name_result: result_dict} - - else: - - logger.log_message(f"Agent '{agent_name}' not found in ai_system.agent_inputs", level=logging.ERROR) - - return {"error": f"Agent '{agent_name}' input configuration not found"} - - else: - - # Multiple agents - execute sequentially - - results = {} - - for agent_name in agent_names: - - single_result = await _execute_custom_agents(ai_system, [agent_name], query) - - results.update(single_result) - - return results - - - - except Exception as e: - - logger.log_message(f"Error in _execute_custom_agents: {str(e)}", level=logging.ERROR) - - return {"error": f"Error executing custom agents: {str(e)}"} - + elif agent_name not in AVAILABLE_AGENTS: + available = list(AVAILABLE_AGENTS.keys()) + raise HTTPException( + status_code=404, + detail=f"Agent '{agent_name}' not found. Available agents: {available}" + ) def _prepare_query_with_context(query: str, session_state: dict) -> str: - """Prepare the query with chat context from previous messages""" - chat_id = session_state.get("chat_id") - if not chat_id: - return query - - # Get chat manager from app state - chat_manager = app.state._session_manager.chat_manager - # Get recent messages - recent_messages = chat_manager.get_recent_chat_history(chat_id, limit=MAX_RECENT_MESSAGES) - # Extract response history - chat_context = chat_manager.extract_response_history(recent_messages) - - # Append context to the query if available - if chat_context: - return f"### Current Query:\n{query}\n\n{chat_context}" - return query - - - def _track_model_usage(session_state: dict, enhanced_query: str, response, processing_time_ms: int): - """Track model usage statistics in the database""" - try: - ai_manager = app.state.get_ai_manager() - - # Get model configuration - model_config = session_state.get("model_config", DEFAULT_MODEL_CONFIG) - model_name = model_config.get("model", DEFAULT_MODEL_CONFIG["model"]) - provider = ai_manager.get_provider_for_model(model_name) - - # Calculate token usage - try: - # Try exact tokenization - prompt_tokens = len(ai_manager.tokenizer.encode(enhanced_query)) - completion_tokens = len(ai_manager.tokenizer.encode(str(response))) - total_tokens = prompt_tokens + completion_tokens - except Exception as token_error: - # Fall back to estimation - logger.log_message(f"Tokenization error: {str(token_error)}", level=logging.WARNING) - prompt_words = len(enhanced_query.split()) - completion_words = len(str(response).split()) - prompt_tokens = int(prompt_words * DEFAULT_TOKEN_RATIO) - completion_tokens = int(completion_words * DEFAULT_TOKEN_RATIO) - total_tokens = prompt_tokens + completion_tokens - - # Calculate cost - cost = ai_manager.calculate_cost(model_name, prompt_tokens, completion_tokens) - - # Save usage to database - ai_manager.save_usage_to_db( - user_id=session_state.get("user_id"), - chat_id=session_state.get("chat_id"), - model_name=model_name, - provider=provider, - prompt_tokens=int(prompt_tokens), - completion_tokens=int(completion_tokens), - total_tokens=int(total_tokens), - query_size=len(enhanced_query), - response_size=len(str(response)), - cost=round(cost, 7), - request_time_ms=processing_time_ms, - is_streaming=False - ) - except Exception as e: - # Log but don't fail the request if usage tracking fails - logger.log_message(f"Failed to track model usage: {str(e)}", level=logging.ERROR) - - - async def _generate_streaming_responses(session_state: dict, query: str, session_lm): - """Generate streaming responses for chat_with_all endpoint""" - overall_start_time = time.time() - total_response = "" - total_inputs = "" - usage_records = [] + try: + # Add chat context from previous messages + enhanced_query = _prepare_query_with_context(query, session_state) + + # Use the session model for this specific request + with dspy.context(lm=session_lm): + try: + # Get the plan + plan_response = await asyncio.wait_for( + asyncio.to_thread(session_state["ai_system"].get_plan, enhanced_query), + timeout=REQUEST_TIMEOUT_SECONDS + ) + + plan_description = format_response_to_markdown( + {"analytical_planner": plan_response}, + dataframe=session_state["current_df"] + ) + + # Check if plan is valid + if plan_description == RESPONSE_ERROR_INVALID_QUERY: + yield json.dumps({ + "agent": "Analytical Planner", + "content": plan_description, + "status": "error" + }) + "\n" + return + + yield json.dumps({ + "agent": "Analytical Planner", + "content": plan_description, + "status": "success" if plan_description else "error" + }) + "\n" + + # Track planner usage + if session_state.get("user_id"): + planner_tokens = _estimate_tokens(ai_manager=app.state.ai_manager, + input_text=enhanced_query, + output_text=plan_description) + + usage_records.append(_create_usage_record( + session_state=session_state, + model_name=session_state.get("model_config", DEFAULT_MODEL_CONFIG)["model"], + prompt_tokens=planner_tokens["prompt"], + completion_tokens=planner_tokens["completion"], + query_size=len(enhanced_query), + response_size=len(plan_description), + processing_time_ms=int((time.time() - overall_start_time) * 1000), + is_streaming=False + )) + + # Execute the plan with well-managed concurrency + async for agent_name, inputs, response in _execute_plan_with_timeout( + session_state["ai_system"], enhanced_query, plan_response): + + if agent_name == "plan_not_found": + yield json.dumps({ + "agent": "Analytical Planner", + "content": "**No plan found**\n\nPlease try again with a different query or try using a different model.", + "status": "error" + }) + "\n" + return + + formatted_response = format_response_to_markdown( + {agent_name: response}, + dataframe=session_state["current_df"] + ) or "No response generated" - - # Add chat context from previous messages - - enhanced_query = _prepare_query_with_context(query, session_state) - - - - # try: - - # Get the plan - planner is now async, so we need to await it - - plan_response = await session_state["ai_system"].get_plan(enhanced_query) - - - - plan_description = format_response_to_markdown( - - {"analytical_planner": plan_response}, - - datasets=session_state["datasets"] - ) - - - - # Check if plan is valid - - if plan_description == RESPONSE_ERROR_INVALID_QUERY: - - yield json.dumps({ - - "agent": "Analytical Planner", - - "content": plan_description, - - "status": "error" - - }) + "\n" - - return - - - - yield json.dumps({ - - "agent": "Analytical Planner", - - "content": plan_description, - - "status": "success" if plan_description else "error" - - }) + "\n" - - - - # Track planner usage - - if session_state.get("user_id"): - - planner_tokens = _estimate_tokens(ai_manager=app.state.ai_manager, - - input_text=enhanced_query, - - output_text=plan_description) - - - - usage_records.append(_create_usage_record( - - session_state=session_state, - - model_name=session_state.get("model_config", DEFAULT_MODEL_CONFIG)["model"], - - prompt_tokens=planner_tokens["prompt"], - - completion_tokens=planner_tokens["completion"], - - query_size=len(enhanced_query), - - response_size=len(plan_description), - - processing_time_ms=int((time.time() - overall_start_time) * 1000), - - is_streaming=False - - )) - - - - logger.log_message(f"Plan response: {plan_response}", level=logging.INFO) - - logger.log_message(f"Plan response type: {type(plan_response)}", level=logging.INFO) - - - - # Check if plan_response is valid - - # if not plan_response or not isinstance(plan_response, dict): - - # yield json.dumps({ - - # "agent": "Analytical Planner", - - # "content": "**Error: Invalid plan response**\n\nResponse: " + str(plan_response), - - # "status": "error" - - # }) + "\n" - - # return - - - - # Execute the plan with well-managed concurrency - - with dspy.context(lm = session_lm): - - # try: - - - - async for agent_name, inputs, response in session_state["ai_system"].execute_plan(enhanced_query, plan_response): - - - - if agent_name == "plan_not_found": - - yield json.dumps({ - - "agent": "Analytical Planner", - - "content": "**No plan found**\n\nPlease try again with a different query or try using a different model.", - - "status": "error" - - }) + "\n" - - return - - - - if agent_name == "plan_not_formated_correctly": - - yield json.dumps({ - - "agent": "Analytical Planner", - - "content": "**Something went wrong with formatting, retry the query!**", - - "status": "error" - - }) + "\n" - - return - - - - - - formatted_response = format_response_to_markdown( - - {agent_name: response}, - - datasets=session_state["datasets"] - ) - - - - yield json.dumps({ - - "agent": agent_name.split("__")[0] if "__" in agent_name else agent_name, - - "content": formatted_response, - - "status": "success" if response else "error" - - }) + "\n" - - - - # Handle agent errors - - if isinstance(response, dict) and "error" in response: - - yield json.dumps({ - - "agent": agent_name, - - "content": f"**Error in {agent_name}**: {response['error']}", - - "status": "error" - - }) + "\n" - - continue # Continue with next agent instead of returning - - - - - - - - if formatted_response == RESPONSE_ERROR_INVALID_QUERY: - - yield json.dumps({ - - "agent": agent_name, - - "content": formatted_response, - - "status": "error" - - }) + "\n" - - continue # Continue with next agent instead of returning - - - - # Send response chunk - - - - - - # Track agent usage for future batch DB write - - if session_state.get("user_id"): - - agent_tokens = _estimate_tokens( - - ai_manager=app.state.ai_manager, - - input_text=str(inputs), - - output_text=str(response) - - ) - - - - # Get appropriate model name for code combiner - - if "code_combiner_agent" in agent_name and "__" in agent_name: - - provider = agent_name.split("__")[1] - - model_name = _get_model_name_for_provider(provider) - - else: - - model_name = session_state.get("model_config", DEFAULT_MODEL_CONFIG)["model"] - - - - usage_records.append(_create_usage_record( - - session_state=session_state, - - model_name=model_name, - - prompt_tokens=agent_tokens["prompt"], - - completion_tokens=agent_tokens["completion"], - - query_size=len(str(inputs)), - - response_size=len(str(response)), - - processing_time_ms=int((time.time() - overall_start_time) * 1000), - - is_streaming=True - - )) - - - - # except asyncio.TimeoutError: - - # yield json.dumps({ - - # "agent": "planner", - - # "content": "The request timed out. Please try a simpler query.", - - # "status": "error" - - # }) + "\n" - - # return - - - - # except Exception as e: - - # logger.log_message(f"Error executing plan: {str(e)}", level=logging.ERROR) - - # yield json.dumps({ - - # "agent": "planner", - - # "content": f"An error occurred while executing the plan: {str(e)}", - - # "status": "error" - - # }) + "\n" - - # return - - - - # except Exception as e: - - # logger.log_message(f"Error in streaming response: {str(e)}", level=logging.ERROR) - - # yield json.dumps({ - - # "agent": "planner", - - # "content": "An error occurred while generating responses. Please try again!" + str(e) + str({k: v for k, v in session_lm.__dict__['kwargs'].items() if k != 'api_key'}), - - # "status": "error" - - # }) + "\n" - - - + if formatted_response == RESPONSE_ERROR_INVALID_QUERY: + yield json.dumps({ + "agent": agent_name, + "content": formatted_response, + "status": "error" + }) + "\n" + return + + # if "code_combiner_agent" in agent_name: + # # logger.log_message(f"[>] Code combiner response: {response}", level=logging.INFO) + # total_response += str(response) if response else "" + # total_inputs += str(inputs) if inputs else "" + + # Send response chunk + yield json.dumps({ + "agent": agent_name.split("__")[0] if "__" in agent_name else agent_name, + "content": formatted_response, + "status": "success" if response else "error" + }) + "\n" + + # Track agent usage for future batch DB write + if session_state.get("user_id"): + agent_tokens = _estimate_tokens( + ai_manager=app.state.ai_manager, + input_text=str(inputs), + output_text=str(response) + ) + + # Get appropriate model name for code combiner + if "code_combiner_agent" in agent_name and "__" in agent_name: + provider = agent_name.split("__")[1] + model_name = _get_model_name_for_provider(provider) + else: + model_name = session_state.get("model_config", DEFAULT_MODEL_CONFIG)["model"] + + usage_records.append(_create_usage_record( + session_state=session_state, + model_name=model_name, + prompt_tokens=agent_tokens["prompt"], + completion_tokens=agent_tokens["completion"], + query_size=len(str(inputs)), + response_size=len(str(response)), + processing_time_ms=int((time.time() - overall_start_time) * 1000), + is_streaming=True + )) + + except asyncio.TimeoutError: + yield json.dumps({ + "agent": "planner", + "content": "The request timed out. Please try a simpler query.", + "status": "error" + }) + "\n" + return + except Exception as e: + logger.log_message(f"Error in streaming response: {str(e)}", level=logging.ERROR) + yield json.dumps({ + "agent": "planner", + "content": "An error occurred while generating responses. Please try again!", + "status": "error" + }) + "\n" + return + + # Batch write usage records to DB + if usage_records and session_state.get("user_id"): + try: + # In a real implementation, you would batch these writes + # For now, we're writing them one by one but could be optimized + ai_manager = app.state.get_ai_manager() + for record in usage_records: + ai_manager.save_usage_to_db(**record) + except Exception as db_error: + logger.log_message(f"Failed to save usage records: {str(db_error)}", level=logging.ERROR) + + except Exception as e: + logger.log_message(f"Streaming response generation failed: {str(e)}", level=logging.ERROR) + yield json.dumps({ + "agent": "planner", + "content": "An error occurred while generating responses. Please try again!", + "status": "error" + }) + "\n" def _estimate_tokens(ai_manager, input_text: str, output_text: str) -> dict: - """Estimate token counts, with fallback for tokenization errors""" - try: - # Try exact tokenization - prompt_tokens = len(ai_manager.tokenizer.encode(input_text)) - completion_tokens = len(ai_manager.tokenizer.encode(output_text)) - except Exception: - # Fall back to estimation - prompt_words = len(input_text.split()) - completion_words = len(output_text.split()) - prompt_tokens = int(prompt_words * DEFAULT_TOKEN_RATIO) - completion_tokens = int(completion_words * DEFAULT_TOKEN_RATIO) - - return { - "prompt": prompt_tokens, - "completion": completion_tokens, - "total": prompt_tokens + completion_tokens - } - - - def _create_usage_record(session_state: dict, model_name: str, prompt_tokens: int, - completion_tokens: int, query_size: int, response_size: int, - processing_time_ms: int, is_streaming: bool) -> dict: - """Create a usage record for the database""" - ai_manager = app.state.get_ai_manager() - provider = ai_manager.get_provider_for_model(model_name) - cost = ai_manager.calculate_cost(model_name, prompt_tokens, completion_tokens) - - return { - "user_id": session_state.get("user_id"), - "chat_id": session_state.get("chat_id"), - "model_name": model_name, - "provider": provider, - "prompt_tokens": int(prompt_tokens), - "completion_tokens": int(completion_tokens), - "total_tokens": int(prompt_tokens + completion_tokens), - "query_size": query_size, - "response_size": response_size, - "cost": round(cost, 7), - "request_time_ms": processing_time_ms, - "is_streaming": is_streaming - } - - - def _get_model_name_for_provider(provider: str) -> str: - """Get the model name for a provider""" - provider_model_map = { - - "openai": "gpt-5.4", - "anthropic": "claude-sonnet-4-6", + "openai": "o3-mini", + "anthropic": "claude-3-7-sonnet-latest", "gemini": "gemini-2.5-pro-preview-03-25" } - return provider_model_map.get(provider, "gpt-5.4") + return provider_model_map.get(provider, "o3-mini") - - - - - -# Add an endpoint to list available agents - -@app.get("/agents", response_model=dict) - -async def list_agents(request: Request, session_id: str = Depends(get_session_id_dependency)): - - """Get all available agents (standard, template, and custom)""" - - session_state = app.state.get_session_state(session_id) - - - +async def _execute_plan_with_timeout(ai_system, enhanced_query, plan_response): + """Execute the plan with timeout handling for each step""" try: - - # Get all available agents from database and session - - available_agents_list = _get_available_agents_list(session_state) - - - - # Categorize agents - - standard_agents = ["preprocessing_agent", "statistical_analytics_agent", "sk_learn_agent", "data_viz_agent"] - - - - # Get template agents from database - - from src.db.init_db import session_factory - - from src.agents.agents import load_all_available_templates_from_db - - - - db_session = session_factory() - - try: - - template_agents_dict = load_all_available_templates_from_db(db_session) - - # template_agents_dict is a dict with template_name as keys - - template_agents = [template_name for template_name in template_agents_dict.keys() - - if template_name not in standard_agents and template_name != 'basic_qa_agent'] - - except Exception as e: - - logger.log_message(f"Error loading template agents in /agents endpoint: {str(e)}", level=logging.ERROR) - - template_agents = [] - - finally: - - db_session.close() - - - - # Get custom agents from session - - custom_agents = [] - - if session_state and "ai_system" in session_state: - - ai_system = session_state["ai_system"] - - if hasattr(ai_system, 'agents'): - - custom_agents = [agent for agent in available_agents_list - - if agent not in standard_agents and agent not in template_agents] - - - - # Ensure template agents are in the available list - - for template_agent in template_agents: - - if template_agent not in available_agents_list: - - available_agents_list.append(template_agent) - - - - return { - - "available_agents": available_agents_list, - - "standard_agents": standard_agents, - - "template_agents": template_agents, - - "custom_agents": custom_agents - - } - + # Use asyncio.create_task to run the execute_plan coroutine + async for agent_name, inputs, response in ai_system.execute_plan(enhanced_query, plan_response): + # Yield results as they come + yield agent_name, inputs, response except Exception as e: - - logger.log_message(f"Error getting agents list: {str(e)}", level=logging.ERROR) - - raise HTTPException(status_code=500, detail=f"Error getting agents list: {str(e)}") + logger.log_message(f"Error executing plan: {str(e)}", level=logging.ERROR) + yield "error", None, {"error": "An error occurred during plan execution"} +# Add an endpoint to list available agents +@app.get("/agents", response_model=dict) +async def list_agents(): + return { + "available_agents": list(AVAILABLE_AGENTS.keys()), + "description": "List of available specialized agents that can be called using @agent_name" + } @app.get("/health", response_model=dict) - async def health(): - return {"message": "API is healthy and running"} - - @app.get("/") - async def index(): - return { - "title": "Welcome to the AI Analytics API", - "message": "Explore our API for advanced analytics and visualization tools designed to empower your data-driven decisions.", - "description": "Utilize our powerful agents and models to gain insights from your data effortlessly.", - "colors": { - "primary": "#007bff", - "secondary": "#6c757d", - "success": "#28a745", - "danger": "#dc3545", - }, - "features": [ - "Real-time data processing", - "Customizable visualizations", - "Seamless integration with various data sources", - "User-friendly interface for easy navigation", - "Custom Analytics", - ], - } - - @app.post("/chat_history_name") - async def chat_history_name(request: dict, session_id: str = Depends(get_session_id_dependency)): - query = request.get("query") - name = None - - - lm = dspy.LM(model="openai/gpt-5-nano", max_tokens=300, temperature=0.5, api_key=os.getenv("OPENAI_API_KEY")) - + lm = dspy.LM(model="gpt-4o-mini", max_tokens=300, temperature=0.5) - with dspy.context(lm=lm): - name = app.state.get_chat_history_name_agent()(query=str(query)) - - return {"name": name.name if name else "New Chat"} - - -@app.post("/deep_analysis_streaming") - -async def deep_analysis_streaming( - - request: DeepAnalysisRequest, - - request_obj: Request, - - session_id: str = Depends(get_session_id_dependency) - -): - - """Perform streaming deep analysis with real-time updates""" - - session_state = app.state.get_session_state(session_id) - - - - try: - - # Extract and validate query parameters - - _update_session_from_query_params(request_obj, session_state) - - - - # Validate dataset - - if session_state["datasets"] is None: - raise HTTPException(status_code=400, detail=RESPONSE_ERROR_NO_DATASET) - - - - # Get user_id from session state (if available) - - user_id = session_state.get("user_id") - - - - # Generate a UUID for this report - - import uuid - - report_uuid = str(uuid.uuid4()) - - - - # Create initial pending report in the database - - try: - - from src.db.init_db import session_factory - - from src.db.schemas.models import DeepAnalysisReport - - - - db_session = session_factory() - - - - try: - - # Create a pending report entry - - new_report = DeepAnalysisReport( - - report_uuid=report_uuid, - - user_id=user_id, - - goal=request.goal, - - status="pending", - - start_time=datetime.now(UTC), - - progress_percentage=0 - - ) - - - - db_session.add(new_report) - - db_session.commit() - - db_session.refresh(new_report) - - - - # Store the report ID in session state for later updates - - session_state["current_deep_analysis_id"] = new_report.report_id - - session_state["current_deep_analysis_uuid"] = report_uuid - - - - except Exception as e: - - logger.log_message(f"Error creating initial deep analysis report: {str(e)}", level=logging.ERROR) - - # Continue even if DB storage fails - - finally: - - db_session.close() - - - - except Exception as e: - - logger.log_message(f"Database operation failed: {str(e)}", level=logging.ERROR) - - # Continue even if DB operation fails - - - - # Get session-specific model - - # session_lm = get_session_lm(session_state) - - session_lm = dspy.LM(model="anthropic/claude-sonnet-4-6", max_tokens=7000, temperature=0.5) - - - - return StreamingResponse( - - _generate_deep_analysis_stream(session_state, request.goal, session_lm, session_id), - - media_type='text/event-stream', - - headers={ - - 'Cache-Control': 'no-cache', - - 'Connection': 'keep-alive', - - 'Content-Type': 'text/event-stream', - - 'Access-Control-Allow-Origin': '*', - - 'X-Accel-Buffering': 'no' - - } - - ) - - - - except HTTPException: - - raise - - except Exception as e: - - logger.log_message(f"Streaming deep analysis failed: {str(e)}", level=logging.ERROR) - - raise HTTPException(status_code=500, detail=f"Streaming deep analysis failed: {str(e)}") - - - -async def _generate_deep_analysis_stream(session_state: dict, goal: str, session_lm, session_id: str): - - """Generate streaming responses for deep analysis""" - - # Track the start time for duration calculation - - start_time = datetime.now(UTC) - - - - try: - - # Get dataset info - datasets = session_state["datasets"] - desc = session_state['description'] - - # Generate dataset info for all datasets - logger.log_message(f"🔍 DEEP ANALYSIS START - datasets type: {type(datasets)}, keys: {list(datasets.keys()) if datasets else 'None'}", level=logging.DEBUG) - - dataset_info = desc - logger.log_message(f"🔍 DEEP ANALYSIS - dataset_info type: {type(dataset_info)}, length: {len(dataset_info) if isinstance(dataset_info, str) else 'N/A'}", level=logging.DEBUG) - logger.log_message(f"🔍 DEEP ANALYSIS - dataset_info content: {dataset_info[:200]}...", level=logging.DEBUG) - - - - # Get report info from session state - - report_id = session_state.get("current_deep_analysis_id") - - report_uuid = session_state.get("current_deep_analysis_uuid") - - user_id = session_state.get("user_id") - - - - # Helper function to update report in database - - async def update_report_in_db(status, progress, step=None, content=None): - - if not report_id: - - return - - - - try: - - from src.db.init_db import session_factory - - from src.db.schemas.models import DeepAnalysisReport - - - - db_session = session_factory() - - - - try: - - report = db_session.query(DeepAnalysisReport).filter(DeepAnalysisReport.report_id == report_id).first() - - - - if report: - - report.status = status - - report.progress_percentage = progress - - - - # Update step-specific fields if provided - - if step == "questions" and content: - - report.deep_questions = content - - elif step == "planning" and content: - - report.deep_plan = content - - elif step == "analysis" and content: - - # For analysis step, we get the full object with multiple fields - - if isinstance(content, dict): - - # Update fields from content if they exist - - if "deep_questions" in content and content["deep_questions"]: - - report.deep_questions = content["deep_questions"] - - if "deep_plan" in content and content["deep_plan"]: - - report.deep_plan = content["deep_plan"] - - if "code" in content and content["code"]: - - report.analysis_code = content["code"] - - if "final_conclusion" in content and content["final_conclusion"]: - - report.final_conclusion = content["final_conclusion"] - - # Also update summary from conclusion - - conclusion = content["final_conclusion"] - - conclusion = conclusion.replace("**Conclusion**", "") - - report.report_summary = conclusion[:200] + "..." if len(conclusion) > 200 else conclusion - - - - # Handle JSON fields - - if "summaries" in content and content["summaries"]: - - report.summaries = json.dumps(content["summaries"]) - - if "plotly_figs" in content and content["plotly_figs"]: - - report.plotly_figures = json.dumps(content["plotly_figs"]) - - if "synthesis" in content and content["synthesis"]: - - report.synthesis = json.dumps(content["synthesis"]) - - - - # For the final step, update the HTML report - - if step == "completed": - - if content: - - report.html_report = content - - else: - - logger.log_message("No HTML content provided for completed step", level=logging.WARNING) - - - - report.end_time = datetime.now(UTC) - - # Ensure start_time is timezone-aware before calculating duration - - if report.start_time.tzinfo is None: - - start_time_utc = report.start_time.replace(tzinfo=UTC) - - else: - - start_time_utc = report.start_time - - report.duration_seconds = int((report.end_time - start_time_utc).total_seconds()) - - - - report.updated_at = datetime.now(UTC) - - db_session.commit() - - - - except Exception as e: - - db_session.rollback() - - logger.log_message(f"Error updating deep analysis report: {str(e)}", level=logging.ERROR) - - finally: - - db_session.close() - - except Exception as e: - - logger.log_message(f"Database operation failed: {str(e)}", level=logging.ERROR) - - - - # Use session model for this request - - with dspy.context(lm=session_lm): - - # Send initial status - - yield json.dumps({ - - "step": "initialization", - - "status": "starting", - - "message": "Initializing deep analysis...", - - "progress": 5 - - }) + "\n" - - - - # Update DB status to running - - await update_report_in_db("running", 5) - - - - # Get deep analyzer - use the correct session_id from the session_state - - logger.log_message(f"Getting deep analyzer for session_id: {session_id}, user_id: {user_id}", level=logging.INFO) - - deep_analyzer = app.state.get_deep_analyzer(session_id) - - - - # Make all datasets available globally for code execution - - for dataset_name, dataset_df in datasets.items(): - globals()[dataset_name] = dataset_df - - # Use the new streaming method and forward all progress updates - final_result = None - - logger.log_message(f"🔍 CALLING DEEP ANALYSIS - goal: {goal[:100]}...", level=logging.DEBUG) - logger.log_message(f"🔍 CALLING DEEP ANALYSIS - dataset_info type: {type(dataset_info)}, length: {len(dataset_info) if isinstance(dataset_info, str) else 'N/A'}", level=logging.DEBUG) - logger.log_message(f"🔍 CALLING DEEP ANALYSIS - session_datasets type: {type(datasets)}, keys: {list(datasets.keys()) if datasets else 'None'}", level=logging.DEBUG) - - async for update in deep_analyzer.execute_deep_analysis_streaming( - goal=goal, - dataset_info=dataset_info, - session_datasets=datasets # Pass all datasets instead of single df - ): - - # Convert the update to the expected format and yield it - - if update.get("step") == "questions" and update.get("status") == "completed": - - # Update DB with questions - - await update_report_in_db("running", update.get("progress", 0), "questions", update.get("content")) - - elif update.get("step") == "planning" and update.get("status") == "completed": - - # Update DB with planning - - await update_report_in_db("running", update.get("progress", 0), "planning", update.get("content")) - - elif update.get("step") == "conclusion" and update.get("status") == "completed": - - # Store the final result for later processing - - final_result = update.get("final_result") - - - - # Convert Plotly figures to JSON format for network transmission - - if final_result: - - import plotly.io - - serialized_return_dict = final_result.copy() - - - - # Convert plotly_figs to JSON format - - if 'plotly_figs' in serialized_return_dict and serialized_return_dict['plotly_figs']: - - json_figs = [] - - for fig_list in serialized_return_dict['plotly_figs']: - - if isinstance(fig_list, list): - - json_fig_list = [] - - for fig in fig_list: - - if hasattr(fig, 'to_json'): # Check if it's a Plotly figure - - json_fig_list.append(plotly.io.to_json(fig)) - - else: - - json_fig_list.append(fig) # Already JSON or other format - - json_figs.append(json_fig_list) - - else: - - # Single figure case - - if hasattr(fig_list, 'to_json'): - - json_figs.append(plotly.io.to_json(fig_list)) - - else: - - json_figs.append(fig_list) - - serialized_return_dict['plotly_figs'] = json_figs - - - - # Update DB with analysis results - - await update_report_in_db("running", update.get("progress", 0), "analysis", serialized_return_dict) - - - - # Generate HTML report using the original final_result with Figure objects - - html_report = None - - try: - - html_report = generate_html_report(final_result) - - except Exception as e: - - logger.log_message(f"Error generating HTML report: {str(e)}", level=logging.ERROR) - - # Continue even if HTML generation fails - - - - # Send the analysis results - - yield json.dumps({ - - "step": "analysis", - - "status": "completed", - - "content": serialized_return_dict, - - "progress": 90 - - }) + "\n" - - - - # Send report generation status - - yield json.dumps({ - - "step": "report", - - "status": "processing", - - "message": "Generating final report...", - - "progress": 95 - - }) + "\n" - - - - # Send final completion - - yield json.dumps({ - - "step": "completed", - - "status": "success", - - "analysis": serialized_return_dict, - - "html_report": html_report, - - "progress": 100 - - }) + "\n" - - - - # Update DB with completed report (with HTML if generated) - - if html_report: - - logger.log_message(f"Saving HTML report to database, length: {len(html_report)}", level=logging.INFO) - - else: - - logger.log_message("No HTML report to save to database", level=logging.WARNING) - - await update_report_in_db("completed", 100, "completed", html_report) - - elif update.get("step") == "error": - - # Forward error directly - - yield json.dumps(update) + "\n" - - await update_report_in_db("failed", 0) - - return - - else: - - # Forward all other progress updates - - yield json.dumps(update) + "\n" - - - - # If we somehow exit the loop without getting a final result, that's an error - - if not final_result: - - yield json.dumps({ - - "step": "error", - - "status": "failed", - - "message": "Deep analysis completed without final result", - - "progress": 0 - - }) + "\n" - - await update_report_in_db("failed", 0) - - - - except Exception as e: - - logger.log_message(f"Error in deep analysis stream: {str(e)}", level=logging.ERROR) - - yield json.dumps({ - - "step": "error", - - "status": "failed", - - "message": f"Deep analysis failed: {str(e)}", - - "progress": 0 - - }) + "\n" - - - - # Update DB with error status - - if 'update_report_in_db' in locals() and session_state.get("current_deep_analysis_id"): - - await update_report_in_db("failed", 0) - - - -@app.post("/deep_analysis/download_report") - -async def download_html_report( - - request: dict, - - session_id: str = Depends(get_session_id_dependency) - -): - - """Download HTML report from previous deep analysis""" - - try: - - analysis_data = request.get("analysis_data") - - if not analysis_data: - - raise HTTPException(status_code=400, detail="No analysis data provided") - - - - # Get report UUID from request if available (for saving to DB) - - report_uuid = request.get("report_uuid") - - session_state = app.state.get_session_state(session_id) - - - - # If no report_uuid in request, try to get it from session state - - if not report_uuid and session_state.get("current_deep_analysis_uuid"): - - report_uuid = session_state.get("current_deep_analysis_uuid") - - - - # Convert JSON-serialized Plotly figures back to Figure objects for HTML generation - - processed_data = analysis_data.copy() - - - - if 'plotly_figs' in processed_data and processed_data['plotly_figs']: - - import plotly.io - - import plotly.graph_objects as go - - - - figure_objects = [] - - for fig_list in processed_data['plotly_figs']: - - if isinstance(fig_list, list): - - fig_obj_list = [] - - for fig_json in fig_list: - - if isinstance(fig_json, str): - - # Convert JSON string back to Figure object - - try: - - fig_obj = plotly.io.from_json(fig_json) - - fig_obj_list.append(fig_obj) - - except Exception as e: - - logger.log_message(f"Error parsing Plotly JSON: {str(e)}", level=logging.WARNING) - - continue - - elif hasattr(fig_json, 'to_html'): - - # Already a Figure object - - fig_obj_list.append(fig_json) - - figure_objects.append(fig_obj_list) - - else: - - # Single figure case - - if isinstance(fig_list, str): - - try: - - fig_obj = plotly.io.from_json(fig_list) - - figure_objects.append(fig_obj) - - except Exception as e: - - logger.log_message(f"Error parsing Plotly JSON: {str(e)}", level=logging.WARNING) - - continue - - elif hasattr(fig_list, 'to_html'): - - figure_objects.append(fig_list) - - - - processed_data['plotly_figs'] = figure_objects - - - - # Generate HTML report - - html_report = generate_html_report(processed_data) - - - - # Save report to database if we have a UUID - - if report_uuid: - - try: - - from src.db.init_db import session_factory - - from src.db.schemas.models import DeepAnalysisReport - - - - db_session = session_factory() - - try: - - # Try to find existing report by UUID - - report = db_session.query(DeepAnalysisReport).filter(DeepAnalysisReport.report_uuid == report_uuid).first() - - - - if report: - - # Update existing report with HTML content - - report.html_report = html_report - - report.updated_at = datetime.now(UTC) - - db_session.commit() - - except Exception as e: - - db_session.rollback() - - finally: - - db_session.close() - - except Exception as e: - - logger.log_message(f"Database operation failed when storing HTML report: {str(e)}", level=logging.ERROR) - - # Continue even if DB storage fails - - - - # Create a filename with timestamp - - timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") - - filename = f"deep_analysis_report_{timestamp}.html" - - - - # Return as downloadable file - - return StreamingResponse( - - iter([html_report.encode('utf-8')]), - - media_type='text/html', - - headers={ - - 'Content-Disposition': f'attachment; filename="{filename}"', - - 'Content-Type': 'text/html; charset=utf-8' - - } - - ) - - - - except Exception as e: - - logger.log_message(f"Failed to generate HTML report: {str(e)}", level=logging.ERROR) - - raise HTTPException(status_code=500, detail=f"Failed to generate report: {str(e)}") - - - - - -# In the section where routers are included, add the session_router - -app.include_router(chat_router) - -app.include_router(analytics_router) - -app.include_router(code_router) - -app.include_router(session_router) - -app.include_router(feedback_router) - -app.include_router(deep_analysis_router) - -app.include_router(templates_router) - -app.include_router(blog_router) - - +# In the section where routers are included, add the session_router +app.include_router(chat_router) +app.include_router(analytics_router) +app.include_router(code_router) +app.include_router(session_router) if __name__ == "__main__": - - port = int(os.environ.get("PORT", 8000)) - - uvicorn.run(app, host="0.0.0.0", port=port) - + uvicorn.run(app, host="0.0.0.0", port=8000) \ No newline at end of file diff --git a/chat_database.db b/chat_database.db new file mode 100644 index 0000000000000000000000000000000000000000..cfb07856ea20f8c9a97b7971cb0789a25ecee31b --- /dev/null +++ b/chat_database.db @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:33b91dbbc8ccfd1c8c384924a5250449ae4092206ad840287cfb6953a2ad220b +size 622592 diff --git a/docs/README.md b/docs/README.md deleted file mode 100644 index 6aaada0d15de98de4ca293ae2ca5cf8c64fef99b..0000000000000000000000000000000000000000 --- a/docs/README.md +++ /dev/null @@ -1,251 +0,0 @@ -# Auto-Analyst Backend Documentation - -This directory contains comprehensive documentation for the Auto-Analyst backend - a sophisticated multi-agent AI platform for data analysis built with FastAPI, DSPy, and modern Python technologies. - -## 📁 Documentation Structure - -### **🏗️ Architecture** (`/architecture/`) -- **[System Architecture](./architecture/architecture.md)** - Comprehensive overview of backend system design, components, and data flow patterns - -### **🚀 Development** (`/development/`) -- **[Development Workflow](./development/development_workflow.md)** - Complete development guide with patterns, best practices, and code organization principles - -### **🔧 System** (`/system/`) -- **[Database Schema](./system/database-schema.md)** - Complete database schema with all tables, relationships, and performance optimization -- **[Shared DataFrame System](./system/shared_dataframe.md)** - Inter-agent data sharing and session management - -### **🌐 API** (`/api/`) -- **[API Endpoints Overview](./api/README.md)** - Main API reference hub -- **[Route Documentation](./api/routes/)** - Detailed endpoint documentation: - - **[Core Routes](./api/routes/session.md)** - File uploads, sessions, authentication - - **[Chat Routes](./api/routes/chats.md)** - Chat and messaging endpoints - - **[Code Routes](./api/routes/code.md)** - Code execution and processing - - **[Analytics Routes](./api/routes/analytics.md)** - Usage analytics and monitoring - - **[Deep Analysis Routes](./api/routes/deep_analysis.md)** - Multi-agent analysis system - - **[Template Routes](./api/routes/templates.md)** - Agent template management - - **[Feedback Routes](./api/routes/feedback.md)** - User feedback and rating system - -### **🐛 Troubleshooting** (`/troubleshooting/`) -- **[Troubleshooting Guide](./troubleshooting/troubleshooting.md)** - Common issues, debugging tools, and solutions - -## 🎯 Backend Overview - -### **Tech Stack** -- **FastAPI** - Modern async Python web framework -- **DSPy** - AI agent orchestration and LLM integration -- **SQLAlchemy** - Database ORM with PostgreSQL/SQLite support -- **Plotly** - Interactive data visualizations -- **Pandas/NumPy** - Data manipulation and analysis -- **Scikit-learn** - Machine learning models -- **Statsmodels** - Statistical analysis - -### **Core Features** -- **Multi-Agent System** - 4+ specialized AI agents for different analysis tasks -- **Template System** - User-customizable agent configurations -- **Deep Analysis** - Multi-step analytical workflows with streaming progress -- **Session Management** - Stateful user sessions with shared data context -- **Code Execution** - Safe Python code execution environment -- **Real-time Streaming** - WebSocket support for live analysis updates - -### **Agent Types** -1. **Data Preprocessing Agent** - Data cleaning and preparation -2. **Statistical Analytics Agent** - Statistical analysis using statsmodels -3. **Machine Learning Agent** - ML modeling with scikit-learn -4. **Data Visualization Agent** - Interactive charts with Plotly -5. **Feature Engineering Agent** (Premium) - Advanced feature creation -6. **Polars Agent** (Premium) - High-performance data processing - -## 🚀 Quick Start Guide - -### **1. Environment Setup** - -```bash -# Navigate to backend directory -cd Auto-Analyst-CS/auto-analyst-backend - -# Create virtual environment -python -m venv venv -source venv/bin/activate # Linux/Mac -venv\Scripts\activate # Windows - -# Install dependencies -pip install -r requirements.txt -``` - -### **2. Environment Configuration** - -Create `.env` file with required variables: - -```env -# Database Configuration -DATABASE_URL=sqlite:///./chat_database.db - -# AI Model Configuration -OPENAI_API_KEY=your-openai-api-key -MODEL_PROVIDER=openai # openai, anthropic, groq, gemini -MODEL_NAME=gpt-4o-mini -TEMPERATURE=0.7 -MAX_TOKENS=6000 - -# Optional: Additional AI Providers -ANTHROPIC_API_KEY=your-anthropic-key -GROQ_API_KEY=your-groq-key -GEMINI_API_KEY=your-gemini-key - -# Security -ADMIN_API_KEY=your-admin-key - -# Application Settings -ENVIRONMENT=development -FRONTEND_URL=http://localhost:3000/ -``` - -### **3. Database Initialization** - -```bash -# Initialize database and default agents -python -c " -from src.db.init_db import init_db -init_db() -print('✅ Database and agents initialized successfully') -" -``` - -### **4. Start Development Server** - -```bash -# Start the FastAPI server -python -m app - -# Or with uvicorn for more control -uvicorn app:app --reload --host 0.0.0.0 --port 8000 -``` - -### **5. Verify Installation** - -- **API Documentation**: `http://localhost:8000/docs` -- **Health Check**: `http://localhost:8000/health` - -## 🔧 Development Workflow - -### **Adding New Agents** - -1. **Define Agent Signature** in `src/agents/agents.py` -2. **Add Configuration** to `agents_config.json` -3. **Register Agent** in loading system -4. **Test Integration** with multi-agent pipeline - -### **Adding New API Endpoints** - -1. **Create Route File** in `src/routes/` -2. **Define Pydantic Models** for request/response -3. **Implement Endpoints** with proper error handling -4. **Register Router** in `app.py` -5. **Update Documentation** - -### **Database Changes** - -1. **Modify Models** in `src/db/schemas/models.py` -2. **Create Migration**: `alembic revision --autogenerate -m "description"` -3. **Apply Migration**: `alembic upgrade head` -4. **Update Documentation** - -## 📊 System Architecture - -### **Request Processing Flow** -``` -HTTP Request → FastAPI Router → Route Handler → Business Logic → -Database/Agent System → AI Model → Response Processing → JSON Response -``` - -### **Agent Execution Flow** -``` -User Query → Session Manager → Agent Selection → Context Preparation → -DSPy Chain → AI Model → Code Generation → Execution → Response Formatting -``` - -### **Deep Analysis Workflow** -``` -Goal Input → Question Generation → Planning → Multi-Agent Execution → -Code Synthesis → Result Compilation → HTML Report Generation -``` - -## 🧪 Testing & Validation - -### **API Testing** -```bash -# Interactive documentation -open http://localhost:8000/docs - -# cURL examples -curl -X GET "http://localhost:8000/health" -curl -X POST "http://localhost:8000/chat/preprocessing_agent" \ - -H "Content-Type: application/json" \ - -d '{"query": "Clean this dataset", "session_id": "test"}' -``` - -### **Agent Testing** -```python -# Test individual agents -from src.agents.agents import preprocessing_agent -import dspy - -# Configure DSPy -lm = dspy.LM('openai/gpt-4o-mini', api_key='your-key') -dspy.configure(lm=lm) - -# Test agent -agent = dspy.ChainOfThought(preprocessing_agent) -result = agent(goal='clean data', dataset='test dataset') -print(result) -``` - -## 🔒 Security & Production - -### **Security Features** -- **Session-based authentication** with secure session management -- **API key protection** for admin endpoints -- **Input validation** using Pydantic models -- **Error handling** with proper HTTP status codes -- **CORS configuration** for frontend integration - -### **Production Considerations** -- **PostgreSQL database** for production deployment -- **Environment variable management** for secrets -- **Logging configuration** for monitoring -- **Rate limiting** for API protection -- **Performance optimization** for large datasets - -## 📈 Monitoring & Analytics - -The backend includes comprehensive analytics for: -- **Usage tracking** - API endpoint usage and performance -- **Model usage** - AI model consumption and costs -- **User analytics** - User behavior and engagement -- **Error monitoring** - System health and error tracking -- **Performance metrics** - Response times and throughput - -## 🤝 Contributing - -1. **Follow coding standards** defined in development workflow -2. **Add comprehensive tests** for new features -3. **Update documentation** for all changes -4. **Use proper error handling** patterns -5. **Submit detailed pull requests** with clear descriptions - ---- - -## 📖 Detailed Documentation - -For specific implementation details, refer to the organized documentation in each subdirectory: - -- **[Getting Started Guide](./getting_started.md)** - Complete setup walkthrough -- **[Architecture Documentation](./architecture/)** - System design and components -- **[Development Guides](./development/)** - Workflow and best practices -- **[API Reference](./api/)** - Complete endpoint documentation -- **[System Documentation](./system/)** - Database and core systems -- **[Troubleshooting](./troubleshooting/)** - Debugging and solutions - ---- - -**Need help?** Check the troubleshooting guide or refer to the comprehensive documentation in each section. \ No newline at end of file diff --git a/docs/api/README.md b/docs/api/README.md deleted file mode 100644 index 1738d93350086855c682ea260a3c027c3d1a4c3c..0000000000000000000000000000000000000000 --- a/docs/api/README.md +++ /dev/null @@ -1,23 +0,0 @@ -# Auto-Analyst Backend API Documentation - -This document is a guide to the backend API endpoints utilized within the Auto-Analyst application. It encompasses a thorough breakdown of various aspects, including the handling of requests, the intricate processes of data transformations, and the structured responses that the API generates. - -The Auto-Analyst application is designed to facilitate seamless interactions and efficient data management, making it essential for users to understand the available endpoints and their functionalities. Each section of this documentation is crafted to provide clarity and insight into how the API operates, ensuring that developers and users alike can effectively leverage its capabilities. - -For more specific details regarding the various functionalities offered by the API, please refer to the following sections, which delve deeper into their respective areas: - -## 📚 Core Documentation - -- **[Getting Started Guide](./getting_started.md)**: Quick start guide for new developers and LLMs to understand the system architecture and get up to speed quickly -- **[System Architecture](./architecture.md)**: Comprehensive overview of the backend system design, components, and data flow patterns -- **[Troubleshooting Guide](./troubleshooting.md)**: Common issues, debugging tools, and solutions for development and deployment problems - -## 🛠️ API Reference - -- **[Core Endpoints](./routes/session.md)**: Review the core endpoints that handle fundamental operations within the application, including data uploads, AI analysis, model settings, and session management. -- **[Analytics Endpoints](./routes/analytics.md)**: Explore the endpoints dedicated to analytics, providing insights into usage statistics, performance metrics, cost analysis, and real-time monitoring. -- **[Chat Endpoints](./routes/chats.md)**: Discover the endpoints that manage chat interactions, enabling users to create, retrieve, and manage chat sessions effectively. -- **[Code Endpoints](./routes/code.md)**: Learn about the endpoints for code execution, editing, fixing, and cleaning operations with advanced AI assistance. -- **[Deep Analysis Endpoints](./routes/deep_analysis.md)**: Comprehensive documentation for the multi-agent deep analysis system, including streaming progress, report management, template integration, and how user's active agents are leveraged for advanced analytical insights. -- **[Feedback Endpoints](./routes/feedback.md)**: Understand the endpoints for managing user feedback on AI-generated messages, including rating systems and model performance tracking. -- **[Templates Endpoints](./routes/templates.md)**: Comprehensive guide to the template system, agent loading, user preferences, and how personalized AI agent configurations work for different users. diff --git a/docs/api/routes/analytics.md b/docs/api/routes/analytics.md deleted file mode 100644 index 8973aeb76cc6494026786a23b7b6b907d88ebbd4..0000000000000000000000000000000000000000 --- a/docs/api/routes/analytics.md +++ /dev/null @@ -1,562 +0,0 @@ -# Analytics Routes Documentation - -These routes provide comprehensive analytics functionality for the Auto-Analyst backend, including dashboard summaries, user analytics, model performance metrics, cost analysis, and system monitoring. - -## Authentication - -All analytics endpoints require admin authentication via API key: - -```python -ADMIN_API_KEY = os.getenv("ADMIN_API_KEY", "default-admin-key-change-me") -``` - -The API key can be provided via: -- **Header:** `X-Admin-API-Key` -- **Query parameter:** `admin_api_key` - ---- - -## Dashboard Endpoints - -### **GET /analytics/dashboard** -Returns comprehensive dashboard data combining usage statistics, model performance, and user activity. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "total_tokens": 123456, - "total_cost": 25.50, - "total_requests": 1000, - "total_users": 50, - "daily_usage": [ - { - "date": "2023-05-01", - "tokens": 5000, - "cost": 1.25, - "requests": 100 - } - ], - "model_usage": [ - { - "model_name": "claude-3-sonnet-20241022", - "tokens": 10000, - "cost": 10.00, - "requests": 200 - } - ], - "top_users": [ - { - "user_id": "123", - "tokens": 5000, - "cost": 5.00, - "requests": 50 - } - ], - "start_date": "2023-04-01", - "end_date": "2023-05-01" -} -``` - -### **WebSocket /analytics/dashboard/realtime** -WebSocket endpoint for real-time dashboard updates. Accepts connections and maintains them for broadcasting live data updates. - ---- - -## User Analytics Endpoints - -### **GET /analytics/users** -Returns user list with usage statistics from the past 7 days. - -**Query Parameters:** -- `limit` (optional): Maximum users to return (default: `100`) -- `offset` (optional): Pagination offset (default: `0`) - -**Response:** -```json -{ - "users": [ - { - "user_id": "123", - "tokens": 5000, - "cost": 5.00, - "requests": 50, - "first_seen": "2023-04-01T12:00:00Z", - "last_seen": "2023-05-01T12:00:00Z" - } - ], - "total": 200, - "limit": 100, - "offset": 0 -} -``` - -### **GET /analytics/users/activity** -Returns daily user activity metrics with new user tracking. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "user_activity": [ - { - "date": "2023-05-01", - "activeUsers": 20, - "newUsers": 5, - "sessions": 30 - } - ] -} -``` - -### **GET /analytics/users/sessions/stats** -Returns session statistics including total users, active users today, average queries per session, and average session time. - -**Response:** -```json -{ - "totalUsers": 500, - "activeToday": 25, - "avgQueriesPerSession": 3.2, - "avgSessionTime": 300 -} -``` - -### **WebSocket /analytics/realtime** -WebSocket endpoint for real-time user analytics updates. - ---- - -## Model Analytics Endpoints - -### **GET /analytics/usage/models** -Returns model usage breakdown with performance metrics. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "model_usage": [ - { - "model_name": "claude-3-sonnet-20241022", - "tokens": 10000, - "cost": 10.00, - "requests": 200, - "avg_response_time": 1.5 - } - ] -} -``` - -### **GET /analytics/models/history** -Returns daily model usage history with trend data. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "model_history": [ - { - "date": "2023-05-01", - "models": [ - { - "name": "claude-3-sonnet-20241022", - "tokens": 5000, - "requests": 100 - } - ] - } - ] -} -``` - -### **GET /analytics/models/metrics** -Returns model performance metrics including success rates and response times. - -**Response:** -```json -{ - "model_metrics": [ - { - "name": "claude-3-sonnet-20241022", - "avg_tokens": 250.5, - "avg_response_time": 1.2, - "success_rate": 0.95 - } - ] -} -``` - ---- - -## Cost Analytics Endpoints - -### **GET /analytics/costs/summary** -Returns cost summary with averages and totals. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "totalCost": 25.50, - "totalTokens": 100000, - "totalRequests": 1000, - "avgDailyCost": 0.85, - "costPerThousandTokens": 0.255, - "daysInPeriod": 30, - "startDate": "2023-04-01", - "endDate": "2023-05-01" -} -``` - -### **GET /analytics/costs/daily** -Returns daily cost breakdown with filled gaps for missing dates. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "daily_costs": [ - { - "date": "2023-05-01", - "cost": 1.25, - "tokens": 5000 - } - ] -} -``` - -### **GET /analytics/costs/models** -Returns cost breakdown by model. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "model_costs": [ - { - "model_name": "claude-3-sonnet-20241022", - "cost": 15.50, - "tokens": 50000, - "requests": 500 - } - ] -} -``` - -### **GET /analytics/costs/projections** -Returns cost projections based on last 30 days usage. - -**Response:** -```json -{ - "nextMonth": 75.00, - "next3Months": 225.00, - "nextYear": 900.00, - "tokensNextMonth": 300000, - "dailyCost": 2.50, - "dailyTokens": 10000, - "baselineDays": 30 -} -``` - -### **GET /analytics/costs/today** -Returns today's cost data. - -**Response:** -```json -{ - "date": "2023-05-01", - "cost": 2.50, - "tokens": 10000, - "requests": 100 -} -``` - ---- - -## Tier Analytics Endpoints - -### **GET /analytics/tiers/usage** -Returns usage data categorized by model tiers with aggregated statistics. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "tier_data": { - "tier_1": { - "name": "Basic", - "credits": 1, - "total_tokens": 50000, - "total_requests": 500, - "total_cost": 5.00, - "avg_tokens_per_query": 100, - "cost_per_1k_tokens": 0.10, - "total_credit_cost": 500, - "cost_per_credit": 0.01, - "models": [...] - } - }, - "period": "30d", - "start_date": "2023-04-01", - "end_date": "2023-05-01" -} -``` - -### **GET /analytics/tiers/projections** -Returns tier-based cost and usage projections. - -**Response:** -```json -{ - "daily_usage": {...}, - "projections": { - "monthly": {...}, - "quarterly": {...}, - "yearly": {...} - }, - "tier_definitions": {...} -} -``` - -### **GET /analytics/tiers/efficiency** -Returns efficiency metrics by tier including cost per credit and tokens per credit. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "efficiency_data": {...}, - "most_efficient_tier": "tier_2", - "best_value_tier": "tier_1", - "period": "30d", - "start_date": "2023-04-01", - "end_date": "2023-05-01" -} -``` - ---- - -## Code Execution Analytics Endpoints - -### **GET /analytics/code-executions/summary** -Returns code execution statistics including success rates and model performance. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "period": "30d", - "start_date": "2023-04-01", - "end_date": "2023-05-01", - "overall_stats": { - "total_executions": 1000, - "successful_executions": 950, - "failed_executions": 50, - "success_rate": 0.95, - "total_users": 100, - "total_chats": 200 - }, - "model_performance": [...], - "failed_agents": [...] -} -``` - -### **GET /analytics/code-executions/detailed** -Returns detailed code execution records with filtering options. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) -- `success_filter` (optional): Filter by success status (boolean) -- `user_id` (optional): Filter by user ID -- `model_name` (optional): Filter by model name -- `limit` (optional): Maximum results (default: `100`) - -**Response:** -```json -{ - "period": "30d", - "start_date": "2023-04-01", - "end_date": "2023-05-01", - "count": 50, - "executions": [...] -} -``` - -### **GET /analytics/code-executions/users** -Returns code execution statistics grouped by user. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) -- `limit` (optional): Maximum users (default: `50`) - -**Response:** -```json -{ - "period": "30d", - "start_date": "2023-04-01", - "end_date": "2023-05-01", - "users": [...] -} -``` - -### **GET /analytics/code-executions/error-analysis** -Returns error analysis with categorized error types and agent failure patterns. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "period": "30d", - "start_date": "2023-04-01", - "end_date": "2023-05-01", - "total_failed_executions": 50, - "error_types": [...], - "error_by_agent": [...] -} -``` - ---- - -## Feedback Analytics Endpoints - -### **GET /analytics/feedback/summary** -Returns feedback summary statistics including rating distributions and trends. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) - -**Response:** -```json -{ - "period": "30d", - "start_date": "2023-04-01", - "end_date": "2023-05-01", - "total_feedback": 500, - "avg_rating": 4.2, - "chats_with_feedback": 200, - "ratings_distribution": [ - {"rating": 1, "count": 10}, - {"rating": 2, "count": 20}, - {"rating": 3, "count": 50}, - {"rating": 4, "count": 200}, - {"rating": 5, "count": 220} - ], - "models_data": [...], - "feedback_trend": [...] -} -``` - -### **GET /analytics/feedback/detailed** -Returns detailed feedback records with filtering and pagination. - -**Query Parameters:** -- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) -- `min_rating` (optional): Minimum rating filter -- `max_rating` (optional): Maximum rating filter -- `model_name` (optional): Filter by model name -- `limit` (optional): Maximum results (default: `100`) -- `offset` (optional): Pagination offset (default: `0`) - -**Response:** -```json -{ - "period": "30d", - "start_date": "2023-04-01", - "end_date": "2023-05-01", - "total": 500, - "count": 100, - "offset": 0, - "limit": 100, - "feedback": [...] -} -``` - ---- - -## Public Endpoints - -### **GET /analytics/public/ticker** -Returns public ticker data for landing page statistics. **No authentication required.** - -**Response:** -```json -{ - "total_signups": 1000, - "total_tokens": 5000000, - "total_requests": 50000, - "last_updated": "2023-05-01T12:00:00Z" -} -``` - ---- - -## Utility Endpoints - -### **GET /analytics/usage/summary** -Returns overall usage summary (legacy endpoint, calls dashboard with 30d period). - -### **GET /analytics/debug/model_usage** -Debug endpoint for testing admin API key validation. - -**Response:** -```json -{ - "status": "success", - "message": "Admin API key validated successfully" -} -``` - ---- - -## Error Categorization - -The system automatically categorizes code execution errors into the following types: - -- **NameError**: Variable or function name not found -- **SyntaxError**: Invalid Python syntax -- **TypeError**: Type-related errors -- **AttributeError**: Attribute access errors -- **IndexError/KeyError**: Index or key access errors -- **ImportError**: Module import errors -- **ValueError**: Invalid values passed to functions -- **OperationError**: Unsupported operations -- **IndentationError**: Python indentation errors -- **PermissionError**: File/system permission errors -- **FileNotFoundError**: File access errors -- **MemoryError**: Memory allocation errors -- **TimeoutError**: Operation timeout errors -- **OtherError**: Uncategorized errors - -## Real-time Updates - -The analytics system supports real-time updates through WebSocket connections: - -- **Dashboard updates**: Broadcasted when new model usage is recorded -- **User activity updates**: Broadcasted for user activity changes -- **Model performance updates**: Broadcasted for model-specific metrics - -All real-time updates are sent as JSON messages with `type` field indicating the update category and `metrics` containing the delta or new values. diff --git a/docs/api/routes/chats.md b/docs/api/routes/chats.md deleted file mode 100644 index 583e18f3a3fc79e00cf595c59d0af97642e3f465..0000000000000000000000000000000000000000 --- a/docs/api/routes/chats.md +++ /dev/null @@ -1,181 +0,0 @@ -### Chat Routes Overview - -These routes handle chat interactions, message processing, user management, and debugging. - ---- - -### **Chat Management** - -#### **1. Create a New Chat** -**Endpoint:** `POST /chats/` -**Description:** Creates a new chat session. -**Request Body:** -```json -{ - "user_id": 123 -} -``` -**Response:** -```json -{ - "chat_id": 456, - "user_id": 123, - "title": "New Chat", - "created_at": "2023-05-01T12:00:00Z" -} -``` - ---- - -#### **2. Retrieve a Chat by ID** -**Endpoint:** `GET /chats/{chat_id}` -**Description:** Fetches a specific chat along with its messages. -**Path Parameter:** `chat_id` (ID of the chat) -**Query Parameter:** `user_id` (Optional for access control) -**Response:** -```json -{ - "chat_id": 456, - "title": "New Chat", - "created_at": "2023-05-01T12:00:00Z", - "user_id": 123, - "messages": [ - { - "message_id": 789, - "chat_id": 456, - "content": "Hello, how can I help?", - "sender": "ai", - "timestamp": "2023-05-01T12:01:00Z" - } - ] -} -``` - ---- - -#### **3. List Recent Chats** -**Endpoint:** `GET /chats/` -**Description:** Retrieves a list of recent chats, optionally filtered by user ID. -**Query Parameters:** -- `user_id` (Optional for filtering by user) -- `limit` (Maximum number of chats, default: 10, max: 100) -- `offset` (For pagination, default: 0) -**Response:** -```json -[ - { - "chat_id": 456, - "user_id": 123, - "title": "New Chat", - "created_at": "2023-05-01T12:00:00Z" - } -] -``` - ---- - -#### **4. Update a Chat** -**Endpoint:** `PUT /chats/{chat_id}` -**Description:** Updates a chat's title or user ID. -**Path Parameter:** `chat_id` (ID of the chat to update) -**Request Body:** -```json -{ - "title": "Updated Chat Title", - "user_id": 123 -} -``` -**Response:** -```json -{ - "chat_id": 456, - "title": "Updated Chat Title", - "created_at": "2023-05-01T12:00:00Z", - "user_id": 123 -} -``` - ---- - -#### **5. Delete a Chat** -**Endpoint:** `DELETE /chats/{chat_id}` -**Description:** Deletes a chat and all its messages while preserving model usage records. -**Path Parameter:** `chat_id` (ID of the chat to delete) -**Query Parameter:** `user_id` (Optional for access control) -**Response:** -```json -{ - "message": "Chat 456 deleted successfully", - "preserved_model_usage": true -} -``` - ---- - -#### **6. Cleanup Empty Chats** -**Endpoint:** `POST /chats/cleanup-empty` -**Description:** Deletes empty chats for a user. -**Request Body:** -```json -{ - "user_id": 123, - "is_admin": false -} -``` -**Response:** -```json -{ - "message": "Deleted 5 empty chats" -} -``` - ---- - -### **Message Management** - -#### **1. Add Message to Chat** -**Endpoint:** `POST /chats/{chat_id}/messages` -**Description:** Adds a message to an existing chat. -**Path Parameter:** `chat_id` (ID of the chat) -**Query Parameter:** `user_id` (Optional for access control) -**Request Body:** -```json -{ - "content": "Hello, I need help with data analysis", - "sender": "user" -} -``` -**Response:** -```json -{ - "message_id": 789, - "chat_id": 456, - "content": "Hello, I need help with data analysis", - "sender": "user", - "timestamp": "2023-05-01T12:01:00Z" -} -``` - ---- - -### **User Management** - -#### **1. Create or Retrieve a User** -**Endpoint:** `POST /chats/users` -**Description:** Creates a new user or retrieves an existing one based on email. -**Request Body:** -```json -{ - "username": "john_doe", - "email": "john@example.com" -} -``` -**Response:** -```json -{ - "user_id": 123, - "username": "john_doe", - "email": "john@example.com", - "created_at": "2023-05-01T12:00:00Z" -} -``` diff --git a/docs/api/routes/code.md b/docs/api/routes/code.md deleted file mode 100644 index 62444e37295620207bdb0be3b70516107e619f0c..0000000000000000000000000000000000000000 --- a/docs/api/routes/code.md +++ /dev/null @@ -1,182 +0,0 @@ -# Code Routes Documentation - -This document describes the API endpoints available for code execution, editing, fixing, and cleaning operations in the Auto-Analyst backend. - -## Base URL - -All code-related endpoints are prefixed with `/code`. - -## Endpoints - -### Execute Code -Executes Python code against the current session's dataframe. - -**Endpoint:** `POST /code/execute` - -**Request Body:** -```json -{ - "code": "string", // Python code to execute - "session_id": "string", // Optional session ID - "message_id": 123 // Optional message ID for tracking -} -``` - -**Response:** -```json -{ - "output": "string", // Execution output - "plotly_outputs": [ // Optional array of plotly outputs - "string" - ] -} -``` - -**Error Responses:** -- `400 Bad Request`: No dataset loaded or no code provided -- `500 Internal Server Error`: Execution error - -### Edit Code -Uses AI to edit code based on user instructions. - -**Endpoint:** `POST /code/edit` - -**Request Body:** -```json -{ - "original_code": "string", // Code to be edited - "user_prompt": "string" // Instructions for editing -} -``` - -**Response:** -```json -{ - "edited_code": "string" // The edited code -} -``` - -**Error Responses:** -- `400 Bad Request`: Missing original code or editing instructions -- `500 Internal Server Error`: Editing error - -### Fix Code -Uses AI to fix code with errors, employing a block-by-block approach with DSPy refinement. - -**Endpoint:** `POST /code/fix` - -**Request Body:** -```json -{ - "code": "string", // Code containing errors - "error": "string" // Error message to fix -} -``` - -**Response:** -```json -{ - "fixed_code": "string" // The fixed code -} -``` - -**Error Responses:** -- `400 Bad Request`: Missing code or error message -- `500 Internal Server Error`: Fixing error - -### Clean Code -Cleans and formats code by organizing imports and ensuring proper code block formatting. - -**Endpoint:** `POST /code/clean-code` - -**Request Body:** -```json -{ - "code": "string" // Code to clean -} -``` - -**Response:** -```json -{ - "cleaned_code": "string" // The cleaned code -} -``` - -**Error Responses:** -- `400 Bad Request`: No code provided -- `500 Internal Server Error`: Cleaning error - -### Get Latest Code -Retrieves the latest code from a specific message. - -**Endpoint:** `POST /code/get-latest-code` - -**Request Body:** -```json -{ - "message_id": 123 // Message ID to retrieve code from -} -``` - -**Response:** -```json -{ - "code": "string" // The retrieved code -} -``` - -**Error Responses:** -- `400 Bad Request`: Missing message ID -- `404 Not Found`: Message not found -- `500 Internal Server Error`: Retrieval error - -## Code Processing Features - -### Import Organization -The code processing system automatically: -- Moves all import statements to the top of the file -- Deduplicates imports -- Sorts imports alphabetically - -### Code Block Management -The system supports code blocks marked with special comments: -- Start marker: `# agent_name code start` -- End marker: `# agent_name code end` - -### Error Handling with DSPy Refinement -When fixing code, the system uses DSPy's refinement mechanism: -- Identifies specific code blocks with errors -- Processes error messages to extract relevant information -- Uses a scoring function to validate fixes -- Employs iterative refinement with up to 3 attempts -- Fixes each block individually while maintaining the overall structure -- Preserves code block markers and relationships - -### Dataset Context -When editing or fixing code, the system provides context about the current dataset including: -- Number of rows and columns -- Column names and data types -- Null value counts -- Sample values for each column - -### Code Execution Safety -The execution system includes safety measures: -- Removes blocking calls like `plt.show()` -- Handles `__main__` block extraction -- Cleans up print statements with unwanted newlines -- Executes code in isolated namespaces - -## Session Management -All endpoints require a valid session ID, which is used to: -- Access the current dataset -- Maintain state between requests -- Track code execution history -- Store execution results for analysis - -## Error Handling -The system provides detailed error messages while maintaining security by: -- Logging errors for debugging -- Returning user-friendly error messages -- Preserving original code in case of processing failures -- Using code scoring to validate fixes before returning results diff --git a/docs/api/routes/deep_analysis.md b/docs/api/routes/deep_analysis.md deleted file mode 100644 index dc1f652a14ef92d21018993eb7561c820c4b6955..0000000000000000000000000000000000000000 --- a/docs/api/routes/deep_analysis.md +++ /dev/null @@ -1,348 +0,0 @@ -# Deep Analysis API Documentation - -## Overview - -The Deep Analysis system provides advanced multi-agent analytical capabilities that automatically generate comprehensive reports based on user goals. The system uses DSPy (Declarative Self-improving Language Programs) to orchestrate multiple AI agents and create detailed analytical insights. - -## Key Features - -- **Multi-Agent Analysis**: Orchestrates multiple specialized agents (preprocessing, statistical analysis, machine learning, visualization) -- **Template Integration**: Uses the user's active templates/agents for analysis -- **Streaming Progress**: Real-time progress updates during analysis execution -- **Report Persistence**: Stores complete analysis reports in database with metadata -- **HTML Export**: Generates downloadable HTML reports with visualizations -- **Credit Tracking**: Monitors token usage, costs, and credits consumed - -## Template Integration - -The deep analysis system integrates with the user's active templates through the agent system: - -1. **Agent Selection**: Uses agents from the user's active template preferences (configured via `/templates` endpoints) -2. **Default Agents**: Falls back to system default agents if user hasn't configured preferences: - - `preprocessing` (both individual and planner variants) - - `statistical_analytics` (both individual and planner variants) - - `sk_learn` (both individual and planner variants) - - `data_viz` (both individual and planner variants) -3. **Template Limits**: Respects the 10-template limit for planner performance optimization -4. **Dynamic Planning**: The planner automatically selects the most appropriate agents based on the analysis goal and available templates - -## Analysis Flow - -The deep analysis process follows these steps: - -1. **Question Generation** (20% progress): Generates 5 targeted analytical questions based on the user's goal -2. **Planning** (40% progress): Creates an optimized execution plan using available agents -3. **Agent Execution** (60% progress): Executes analysis using user's active templates -4. **Code Synthesis** (80% progress): Combines and optimizes code from all agents -5. **Code Execution** (85% progress): Runs the synthesized analysis code -6. **Synthesis** (90% progress): Synthesizes results into coherent insights -7. **Conclusion** (100% progress): Generates final conclusions and recommendations - ---- - -## API Endpoints - -### Create Deep Analysis Report - -**POST** `/deep_analysis/reports` - -Creates a new deep analysis report in the database. - -**Request Body:** -```json -{ - "report_uuid": "string", - "user_id": 123, - "goal": "Analyze customer churn patterns", - "status": "completed", - "deep_questions": "1. What factors...\n2. How does...", - "deep_plan": "{\n \"@preprocessing\": {\n \"create\": [...],\n \"use\": [...],\n \"instruction\": \"...\"\n }\n}", - "summaries": ["Agent summary 1", "Agent summary 2"], - "analysis_code": "import pandas as pd\n# Analysis code...", - "plotly_figures": [{"data": [...], "layout": {...}}], - "synthesis": ["Synthesis result 1"], - "final_conclusion": "## Conclusion\nThe analysis reveals...", - "html_report": "...", - "report_summary": "Brief summary of findings", - "progress_percentage": 100, - "duration_seconds": 120, - "credits_consumed": 5, - "error_message": null, - "model_provider": "anthropic", - "model_name": "claude-sonnet-4-20250514", - "total_tokens_used": 15000, - "estimated_cost": 0.25, - "steps_completed": ["questions", "planning", "execution", "synthesis", "conclusion"] -} -``` - -**Response:** -```json -{ - "report_id": 1, - "report_uuid": "uuid-string", - "user_id": 123, - "goal": "Analyze customer churn patterns", - "status": "completed", - "start_time": "2024-01-01T12:00:00Z", - "end_time": "2024-01-01T12:02:00Z", - "duration_seconds": 120, - "report_summary": "Brief summary of findings", - "created_at": "2024-01-01T12:02:00Z", - "updated_at": "2024-01-01T12:02:00Z" -} -``` - -### Get Deep Analysis Reports - -**GET** `/deep_analysis/reports` - -Retrieves a list of deep analysis reports with optional filtering. - -**Query Parameters:** -- `user_id` (optional): Filter by user ID -- `limit` (optional): Number of reports to return (1-100, default: 10) -- `offset` (optional): Number of reports to skip (default: 0) -- `status` (optional): Filter by status ("pending", "running", "completed", "failed") - -**Response:** -```json -[ - { - "report_id": 1, - "report_uuid": "uuid-string", - "user_id": 123, - "goal": "Analyze customer churn patterns", - "status": "completed", - "start_time": "2024-01-01T12:00:00Z", - "end_time": "2024-01-01T12:02:00Z", - "duration_seconds": 120, - "report_summary": "Brief summary of findings", - "created_at": "2024-01-01T12:02:00Z", - "updated_at": "2024-01-01T12:02:00Z" - } -] -``` - -### Get User Historical Reports - -**GET** `/deep_analysis/reports/user_historical` - -Retrieves all historical deep analysis reports for a specific user. - -**Query Parameters:** -- `user_id`: User ID (required) -- `limit` (optional): Number of reports to return (1-100, default: 50) - -### Get Report by ID - -**GET** `/deep_analysis/reports/{report_id}` - -Retrieves a complete deep analysis report by ID. - -**Query Parameters:** -- `user_id` (optional): Ensures report belongs to specified user - -**Response:** -```json -{ - "report_id": 1, - "report_uuid": "uuid-string", - "user_id": 123, - "goal": "Analyze customer churn patterns", - "status": "completed", - "start_time": "2024-01-01T12:00:00Z", - "end_time": "2024-01-01T12:02:00Z", - "duration_seconds": 120, - "deep_questions": "1. What factors contribute to churn?\n2. How does churn vary by segment?", - "deep_plan": "{\n \"@preprocessing\": {...},\n \"@statistical_analytics\": {...}\n}", - "summaries": ["Agent performed data cleaning...", "Statistical analysis revealed..."], - "analysis_code": "import pandas as pd\n# Complete analysis code", - "plotly_figures": [{"data": [...], "layout": {...}}], - "synthesis": ["The analysis shows clear patterns..."], - "final_conclusion": "## Conclusion\nCustomer churn is primarily driven by...", - "html_report": "...", - "report_summary": "Analysis of customer churn patterns reveals...", - "progress_percentage": 100, - "credits_consumed": 5, - "error_message": null, - "model_provider": "anthropic", - "model_name": "claude-sonnet-4-20250514", - "total_tokens_used": 15000, - "estimated_cost": 0.25, - "steps_completed": ["questions", "planning", "execution", "synthesis", "conclusion"], - "created_at": "2024-01-01T12:02:00Z", - "updated_at": "2024-01-01T12:02:00Z" -} -``` - -### Get Report by UUID - -**GET** `/deep_analysis/reports/uuid/{report_uuid}` - -Retrieves a complete deep analysis report by UUID. Same response format as get by ID. - -### Delete Report - -**DELETE** `/deep_analysis/reports/{report_id}` - -Deletes a deep analysis report. - -**Query Parameters:** -- `user_id` (optional): Ensures report belongs to specified user - -**Response:** -```json -{ - "message": "Report 1 deleted successfully" -} -``` - -### Update Report Status - -**PUT** `/deep_analysis/reports/{report_id}/status` - -Updates the status of a deep analysis report. - -**Request Body:** -```json -{ - "status": "completed" -} -``` - -**Valid Status Values:** -- `pending`: Analysis queued but not started -- `running`: Analysis in progress -- `completed`: Analysis finished successfully -- `failed`: Analysis encountered errors - -### Get HTML Report - -**GET** `/deep_analysis/reports/uuid/{report_uuid}/html` - -Retrieves only the HTML report content for a specific analysis. - -**Query Parameters:** -- `user_id` (optional): Ensures report belongs to specified user - -**Response:** -```json -{ - "html_report": "...", - "filename": "deep_analysis_report_20240101_120200.html" -} -``` - -### Download HTML Report - -**POST** `/deep_analysis/download_from_db/{report_uuid}` - -Downloads the HTML report as a file attachment. - -**Query Parameters:** -- `user_id` (optional): Ensures report belongs to specified user - -**Response:** -- Content-Type: `text/html; charset=utf-8` -- Content-Disposition: `attachment; filename="deep_analysis_report_TIMESTAMP.html"` - ---- - -## Deep Analysis Module Architecture - -### DSPy Signatures - -The system uses several DSPy signatures for different analysis phases: - -#### 1. `deep_questions` -Generates 5 targeted analytical questions based on the user's goal and dataset structure. - -#### 2. `deep_planner` -Creates an optimized execution plan using the user's active templates/agents. The planner: -- Verifies feasibility using available datasets and agent descriptions -- Batches similar questions per agent call for efficiency -- Reuses outputs across questions to minimize agent calls -- Defines clear variable flow and dependencies between agents - -#### 3. `deep_code_synthesizer` -Combines and optimizes code from multiple agents: -- Fixes errors and inconsistencies between agent outputs -- Ensures proper data flow and type handling -- Converts all visualizations to Plotly format -- Adds comprehensive error handling and validation - -#### 4. `deep_synthesizer` -Synthesizes analysis results into coherent insights and findings. - -#### 5. `final_conclusion` -Generates final conclusions and strategic recommendations based on all analysis results. - -### Streaming Analysis - -The `execute_deep_analysis_streaming` method provides real-time progress updates: - -```python -async for update in deep_analysis.execute_deep_analysis_streaming(goal, dataset_info, session_df): - if update["step"] == "questions": - # Handle questions generation progress - elif update["step"] == "planning": - # Handle planning progress - elif update["step"] == "agent_execution": - # Handle agent execution progress - # ... handle other steps -``` - -### Integration with User Templates - -The deep analysis system integrates with user templates in several ways: - -1. **Agent Discovery**: Retrieves user's active template preferences from the database -2. **Dynamic Planning**: The planner uses available agents to create optimal execution plans -3. **Template Validation**: Ensures all referenced agents exist in the user's active templates -4. **Fallback Handling**: Uses default agents if user preferences are incomplete -5. **Performance Optimization**: Respects template limits for efficient execution - -### Error Handling - -The system includes comprehensive error handling: - -- **Code Execution Errors**: Automatically attempts to fix and retry failed code -- **Template Missing**: Falls back to default agents if user templates are unavailable -- **Timeout Protection**: Includes timeouts for long-running operations -- **Memory Management**: Handles large datasets and visualization efficiently -- **Unicode Handling**: Cleans problematic characters that might cause encoding issues - -### Visualization Integration - -All visualizations are standardized to Plotly format: -- Consistent styling and color schemes -- Interactive features (zoom, pan, hover) -- Accessibility compliance (colorblind-friendly palettes) -- Export capabilities for reports -- Responsive design for different screen sizes - ---- - -## Frontend Integration - -The deep analysis system includes React components for: - -- **DeepAnalysisSidebar**: Main interface for starting and managing analyses -- **NewAnalysisForm**: Form for initiating new deep analyses -- **CurrentAnalysisView**: Real-time progress tracking during analysis -- **HistoryView**: Browse and access historical analysis reports -- **AnalysisStep**: Individual step progress visualization - -The frontend integrates with the streaming API to provide real-time feedback and uses the user's active template configuration for personalized analysis capabilities. - -## Credit and Cost Tracking - -The system tracks detailed usage metrics: -- **Credits Consumed**: Number of credits deducted from user account -- **Token Usage**: Total tokens used across all model calls -- **Estimated Cost**: Dollar cost estimate based on model pricing -- **Model Information**: Provider and model name used for analysis -- **Execution Time**: Duration of analysis for performance monitoring - -This information helps users understand resource consumption and optimize their analysis strategies. \ No newline at end of file diff --git a/docs/api/routes/feedback.md b/docs/api/routes/feedback.md deleted file mode 100644 index 45dda6933cf87e62f0b26107e580cfc7fd5ce103..0000000000000000000000000000000000000000 --- a/docs/api/routes/feedback.md +++ /dev/null @@ -1,153 +0,0 @@ -# Feedback Routes Documentation - -This document describes the API endpoints available for managing user feedback on AI-generated messages in the Auto-Analyst backend. - -## Base URL - -All feedback-related endpoints are prefixed with `/feedback`. - -## Endpoints - -### Create or Update Message Feedback -Creates new feedback or updates existing feedback for a specific message. - -**Endpoint:** `POST /feedback/message/{message_id}` - -**Path Parameters:** -- `message_id`: ID of the message to provide feedback for - -**Request Body:** -```json -{ - "rating": 5, // Required: Star rating (1-5) - "model_name": "gpt-4o-mini", // Optional: Model used for the message - "model_provider": "openai", // Optional: Provider of the model - "temperature": 0.7, // Optional: Temperature setting - "max_tokens": 6000 // Optional: Max tokens setting -} -``` - -**Response:** -```json -{ - "feedback_id": 123, - "message_id": 456, - "rating": 5, - "feedback_comment": null, - "model_name": "gpt-4o-mini", - "model_provider": "openai", - "temperature": 0.7, - "max_tokens": 6000, - "created_at": "2023-05-01T12:00:00Z", - "updated_at": "2023-05-01T12:00:00Z" -} -``` - -**Error Responses:** -- `404 Not Found`: Message with specified ID not found -- `500 Internal Server Error`: Failed to create/update feedback - -### Get Message Feedback -Retrieves feedback for a specific message. - -**Endpoint:** `GET /feedback/message/{message_id}` - -**Path Parameters:** -- `message_id`: ID of the message to get feedback for - -**Response:** -```json -{ - "feedback_id": 123, - "message_id": 456, - "rating": 5, - "feedback_comment": null, - "model_name": "gpt-4o-mini", - "model_provider": "openai", - "temperature": 0.7, - "max_tokens": 6000, - "created_at": "2023-05-01T12:00:00Z", - "updated_at": "2023-05-01T12:00:00Z" -} -``` - -**Error Responses:** -- `404 Not Found`: No feedback found for the specified message -- `500 Internal Server Error`: Failed to retrieve feedback - -### Get Chat Feedback -Retrieves all feedback for messages in a specific chat. - -**Endpoint:** `GET /feedback/chat/{chat_id}` - -**Path Parameters:** -- `chat_id`: ID of the chat to get feedback for - -**Response:** -```json -[ - { - "feedback_id": 123, - "message_id": 456, - "rating": 5, - "feedback_comment": null, - "model_name": "gpt-4o-mini", - "model_provider": "openai", - "temperature": 0.7, - "max_tokens": 6000, - "created_at": "2023-05-01T12:00:00Z", - "updated_at": "2023-05-01T12:00:00Z" - } -] -``` - -**Note:** Returns an empty array if no feedback exists for the chat. - -**Error Responses:** -- `500 Internal Server Error`: Failed to retrieve chat feedback - -## Feedback Features - -### Rating System -- **Scale:** 1-5 star rating system -- **Required:** Rating is the only required field for feedback -- **Purpose:** Allows users to rate the quality of AI responses - -### Model Context Tracking -The system optionally tracks: -- **Model Name:** The specific AI model used (e.g., "gpt-4o-mini") -- **Model Provider:** The provider of the model (e.g., "openai", "anthropic") -- **Temperature:** The creativity/randomness setting used -- **Max Tokens:** The maximum response length setting - -### Update Behavior -- **Upsert Operation:** The POST endpoint either creates new feedback or updates existing feedback -- **Partial Updates:** When updating, only provided fields are modified -- **Timestamp Tracking:** Both creation and update timestamps are maintained - -## Data Management - -### Database Operations -- **Atomic Operations:** Feedback creation/updates are handled in database transactions -- **Referential Integrity:** Feedback is linked to specific messages via foreign keys -- **Soft Handling:** Missing optional fields are handled gracefully - -### Error Handling -- **Comprehensive Logging:** All operations are logged for debugging -- **User-Friendly Messages:** Error responses provide clear information -- **Transaction Safety:** Failed operations are rolled back to maintain data consistency - -## Usage Patterns - -### Typical Workflow -1. User receives an AI-generated message -2. User provides rating (1-5 stars) via the frontend -3. Frontend calls `POST /feedback/message/{message_id}` with rating and model context -4. System stores or updates the feedback -5. Feedback can be retrieved later for analytics or user review - -### Analytics Integration -Feedback data is used by the analytics system to: -- Track model performance across different configurations -- Identify patterns in user satisfaction -- Generate insights for model optimization \ No newline at end of file diff --git a/docs/api/routes/session.md b/docs/api/routes/session.md deleted file mode 100644 index 18d1570c992c65c0353749350f6570b8c40beaaa..0000000000000000000000000000000000000000 --- a/docs/api/routes/session.md +++ /dev/null @@ -1,273 +0,0 @@ -# **Auto-Analyst API Documentation** - -The core application routes are designed to manage the data and AI analysis capabilities of the Auto-Analyst application. - -## **1. Core Application Routes** -### **Data Management** - -#### **POST /upload_dataframe** -Uploads a CSV dataset for analysis. -**Request:** -- `file`: CSV file -- `name`: Dataset name -- `description`: Dataset description -**Headers:** -- `X-Force-Refresh`: "true" (optional) - Forces session reset before upload -**Response:** -```json -{ "message": "Dataframe uploaded successfully", "session_id": "abc123" } -``` - -#### **POST /upload_excel** -Uploads an Excel file with a specific sheet for analysis. -**Request:** -- `file`: Excel file -- `name`: Dataset name -- `description`: Dataset description -- `sheet_name`: Name of the Excel sheet to use -**Headers:** -- `X-Force-Refresh`: "true" (optional) - Forces session reset before upload -**Response:** -```json -{ "message": "Excel file processed successfully", "session_id": "abc123", "sheet": "Sheet1" } -``` - -#### **POST /api/excel-sheets** -Gets the list of sheet names from an Excel file. -**Request:** -- `file`: Excel file -**Response:** -```json -{ "sheets": ["Sheet1", "Sheet2", "Data"] } -``` - -#### **GET /api/default-dataset** -Gets the default dataset. -**Response:** -```json -{ - "headers": ["column1", "column2", ...], - "rows": [[val1, val2, ...], ...], - "name": "Housing Dataset", - "description": "A comprehensive dataset containing housing information..." -} -``` - -#### **POST /reset-session** -Resets session to default dataset. -**Request Body:** -```json -{ - "name": "optional name", - "description": "optional description", - "preserveModelSettings": false -} -``` -**Response:** -```json -{ - "message": "Session reset to default dataset", - "session_id": "abc123", - "dataset": "Housing.csv" -} -``` - -#### **GET /api/preview-csv** / **POST /api/preview-csv** -Preview the current dataset in the session. -**Response:** -```json -{ - "headers": ["column1", "column2", ...], - "rows": [[val1, val2, ...], ...], - "name": "Dataset Name", - "description": "Dataset description..." -} -``` - ---- - -### **2. AI Analysis** - -#### **POST /chat/{agent_name}** -Processes a query using a specific AI agent. -**Path Parameters:** `agent_name` -**Request Body:** -```json -{ "query": "Analyze the relationship between price and size" } -``` -**Query Parameters:** `user_id` (optional), `chat_id` (optional) -**Response:** -```json -{ - "agent_name": "data_viz_agent", - "query": "Analyze the relationship between price and size", - "response": "# Analysis\n\nThere appears to be a strong positive correlation...", - "session_id": "abc123" -} -``` - -#### **POST /chat** -Processes a query using multiple AI agents with streaming responses. -**Request Body:** -```json -{ "query": "Analyze the housing data" } -``` -**Query Parameters:** `user_id` (optional), `chat_id` (optional) -**Response:** *Streaming JSON objects:* -```json -{"agent": "data_viz_agent", "content": "# Visualization\n\n...", "status": "success"} -{"agent": "statistical_analytics_agent", "content": "# Statistical Analysis\n\n...", "status": "success"} -``` - -#### **POST /chat_history_name** -Generates a name for a chat based on the query. -**Request Body:** -```json -{ "query": "Analyze sales data for Q4" } -``` -**Response:** -```json -{ "name": "Chat about sales data analysis" } -``` - -#### **GET /agents** -Lists available AI agents. -**Response:** -```json -{ - "available_agents": ["data_viz_agent", "sk_learn_agent", "statistical_analytics_agent", "preprocessing_agent"], - "standard_agents": ["preprocessing_agent", "statistical_analytics_agent", "sk_learn_agent", "data_viz_agent"], - "template_agents": ["custom_template_1", "custom_template_2"], - "custom_agents": [] -} -``` - ---- - -### **3. Deep Analysis** - -#### **POST /deep_analysis_streaming** -Performs comprehensive deep analysis with real-time streaming updates. -**Request Body:** -```json -{ "goal": "Perform comprehensive analysis of the sales data" } -``` -**Query Parameters:** `user_id` (optional), `chat_id` (optional) -**Response:** *Streaming JSON objects with progress updates* - -#### **POST /deep_analysis/download_report** -Downloads an HTML report from deep analysis results. -**Request Body:** -```json -{ - "analysis_data": { ... }, - "report_uuid": "optional-uuid" -} -``` -**Response:** HTML file download - ---- - -### **4. Model Settings** - -#### **GET /api/model-settings** -Fetches current model settings. -**Response:** -```json -{ - "provider": "openai", - "model": "gpt-4o-mini", - "hasCustomKey": true, - "temperature": 1.0, - "maxTokens": 6000 -} -``` - -#### **POST /settings/model** -Updates model settings. -**Request Body:** -```json -{ - "provider": "openai", - "model": "gpt-4", - "api_key": "sk-...", - "temperature": 0.7, - "max_tokens": 8000 -} -``` -**Response:** -```json -{ "message": "Model settings updated successfully" } -``` - ---- - -### **5. Session Management** - -#### **GET /api/session-info** -Gets information about the current session. -**Response:** -```json -{ - "session_id": "abc123", - "dataset_name": "Housing Dataset", - "dataset_description": "...", - "model_config": { ... } -} -``` - -#### **POST /set-message-info** -Associates message tracking information with the session. -**Request Body:** -```json -{ - "chat_id": 123, - "message_id": 456, - "user_id": 789 -} -``` - -#### **POST /create-dataset-description** -Creates an AI-generated description for a dataset. -**Request Body:** -```json -{ - "df_preview": "column1,column2\nvalue1,value2\n...", - "name": "Dataset Name" -} -``` - ---- - -### **6. System Endpoints** - -#### **GET /** -Returns API welcome information and feature list. - -#### **GET /health** -Health check endpoint. -**Response:** -```json -{ "message": "API is healthy and running" } -``` - ---- - ---- - -### **7. Authentication & Session Management** -- **Session ID Sources:** - - Query parameter: `session_id` - - Header: `X-Session-ID` - - Auto-generated if not provided -- **Session State Includes:** - - Current dataset - - AI system instance - - Model configuration - - User and chat associations - -### **9. Error Handling** -- Comprehensive error handling with appropriate HTTP status codes -- Detailed error messages for debugging -- Fallback encoding support for CSV files (UTF-8, unicode_escape, ISO-8859-1) -- Session state preservation during errors diff --git a/docs/api/routes/templates.md b/docs/api/routes/templates.md deleted file mode 100644 index 29803b79585c0ba6baf4090364ede7d737f2b692..0000000000000000000000000000000000000000 --- a/docs/api/routes/templates.md +++ /dev/null @@ -1,363 +0,0 @@ -# Templates and Agent Loading Documentation - -This document describes how the Auto-Analyst template system works, including agent loading, user preferences, and template management. - -## Overview - -The Auto-Analyst system uses a flexible template-based approach for managing AI agents. Templates define specialized agents with specific capabilities, and users can customize which agents are available for their analysis workflows. - -## Template System Architecture - -### Template Types - -Templates come in different **variant types** that determine how they can be used: - -- **`individual`**: Templates available for single-agent queries (e.g., `@preprocessing_agent`) -- **`planner`**: Templates available for multi-agent planning workflows -- **`both`**: Templates available in both individual and planner contexts - -### Default Agents - -The system includes four core default agents that are **enabled by default** for all users: - -**For Individual Use:** -- `preprocessing_agent`: Data cleaning and preprocessing -- `statistical_analytics_agent`: Statistical analysis and insights -- `sk_learn_agent`: Machine learning with scikit-learn -- `data_viz_agent`: Data visualization with Plotly - -**For Planner Use:** -- `planner_preprocessing_agent`: Planning version of preprocessing agent -- `planner_statistical_analytics_agent`: Planning version of statistical agent -- `planner_sk_learn_agent`: Planning version of ML agent -- `planner_data_viz_agent`: Planning version of visualization agent - -## Template Management Endpoints - -### Get All Templates - -**Endpoint:** `GET /templates/` - -**Query Parameters:** -- `variant_type`: Filter by `"individual"`, `"planner"`, or `"all"` (default: `"all"`) - -**Response:** -```json -[ - { - "template_id": 1, - "template_name": "preprocessing_agent", - "display_name": "Data Preprocessing Agent", - "description": "Handles data cleaning, missing values, and preprocessing tasks", - "prompt_template": "You are a data preprocessing specialist...", - "template_category": "Data Processing", - "icon_url": "/icons/templates/preprocessing_agent.svg", - "is_premium_only": false, - "is_active": true, - "usage_count": 12, - "created_at": "2023-05-01T12:00:00Z", - "updated_at": "2023-05-01T12:00:00Z" - } -] -``` - -### Get Templates by Category - -**Endpoint:** `GET /templates/categories` - -**Query Parameters:** -- `variant_type`: Filter by `"individual"`, `"planner"`, or `"all"` (default: `"individual"`) - -**Response:** -```json -[ - { - "category": "Data Processing", - "templates": [ - { - "agent_id": 1, - "agent_name": "preprocessing_agent", - "display_name": "Data Preprocessing Agent", - "description": "Handles data cleaning and preprocessing", - "icon_url": "/icons/templates/preprocessing_agent.svg", - "usage_count": 1234 - } - ] - } -] -``` - -### Get Template by ID - -**Endpoint:** `GET /templates/template/{template_id}` - -**Response:** -```json -{ - "template_id": 1, - "template_name": "preprocessing_agent", - "display_name": "Data Preprocessing Agent", - "description": "Handles data cleaning, missing values, and preprocessing tasks", - "prompt_template": "You are a data preprocessing specialist...", - "template_category": "Data Processing", - "icon_url": "/icons/templates/preprocessing_agent.svg", - "is_premium_only": false, - "is_active": true, - "usage_count": 1234, - "created_at": "2023-05-01T12:00:00Z", - "updated_at": "2023-05-01T12:00:00Z" -} -``` - -### Get Template Categories List - -**Endpoint:** `GET /templates/categories/list` - -**Response:** -```json -{ - "categories": [ - "Data Processing", - "Machine Learning", - "Visualization", - "Statistics" - ] -} -``` - -### Get Templates by Specific Category - -**Endpoint:** `GET /templates/category/{category}` - -**Path Parameters:** -- `category`: Name of the category to filter by - -**Response:** -```json -[ - { - "template_id": 1, - "template_name": "preprocessing_agent", - "display_name": "Data Preprocessing Agent", - "description": "Handles data cleaning and preprocessing", - "template_category": "Data Processing", - "icon_url": "/icons/templates/preprocessing_agent.svg", - "usage_count": 1234 - } -] -``` - -## User Template Preferences - -### How Agent Loading Works for Users - -1. **Default Behavior**: New users automatically have the 4 core default agents enabled -2. **Custom Preferences**: Users can enable/disable additional templates through preferences -3. **Variant-Specific**: Individual and planner variants are managed separately -4. **Usage Tracking**: System tracks which templates users actually use - -### Get User Template Preferences - -**Endpoint:** `GET /templates/user/{user_id}` - -**Query Parameters:** -- `variant_type`: Filter by `"individual"`, `"planner"`, or `"all"` (default: `"planner"`) - -**Response:** -```json -[ - { - "template_id": 1, - "template_name": "preprocessing_agent", - "display_name": "Data Preprocessing Agent", - "description": "Handles data cleaning and preprocessing", - "template_category": "Data Processing", - "icon_url": "/icons/templates/preprocessing_agent.svg", - "is_premium_only": false, - "is_active": true, - "is_enabled": true, - "usage_count": 15, - "last_used_at": "2023-05-01T12:00:00Z", - "created_at": "2023-04-01T12:00:00Z", - "updated_at": "2023-05-01T12:00:00Z" - } -] -``` - -### Get Only Enabled Templates - -**Endpoint:** `GET /templates/user/{user_id}/enabled` - -Returns only templates that are currently enabled for the user. - -### Get Enabled Templates for Planner - -**Endpoint:** `GET /templates/user/{user_id}/enabled/planner` - -Returns enabled planner templates with the following restrictions: -- **Maximum 10 templates** for planner use -- **Sorted by usage** (most used first) -- **Only planner variants** (`planner` or `both` types) - -## Template Preference Management - -### Toggle Single Template - -**Endpoint:** `POST /templates/user/{user_id}/template/{template_id}/toggle` - -**Request Body:** -```json -{ - "is_enabled": true -} -``` - -**Restrictions:** -- Cannot disable all templates (at least 1 must remain enabled) -- Cannot enable more than 10 templates for planner use - -### Bulk Toggle Templates - -**Endpoint:** `POST /templates/user/{user_id}/bulk-toggle` - -**Request Body:** -```json -{ - "preferences": [ - { - "template_id": 1, - "is_enabled": true - }, - { - "template_id": 2, - "is_enabled": false - } - ] -} -``` - -**Response:** -```json -{ - "results": [ - { - "template_id": 1, - "success": true, - "message": "Template enabled successfully", - "is_enabled": true - } - ] -} -``` - -## Template Categories and Icons - -### Available Categories - -Templates are organized into categories such as: -- **Data Processing**: Preprocessing, cleaning, feature engineering -- **Machine Learning**: Various ML frameworks and algorithms -- **Visualization**: Plotting and chart generation -- **Statistics**: Statistical analysis and modeling -- **Custom**: User or organization-specific templates - -### Icon System - -Templates include visual icons stored in `/public/icons/templates/`: - -**Core Agent Icons:** -- `preprocessing_agent.svg`: Data preprocessing -- `sk_learn_agent.svg`: Machine learning -- `matplotlib_agent.png`: Plotting with matplotlib -- `polars_agent.svg`: Data manipulation with Polars - -**Library-Specific Icons:** -- `numpy.svg`, `scipy.png`: Scientific computing -- `plotly.svg`, `seaborn.svg`: Advanced visualization -- `lightgbm.png`, `xgboost.png`: Gradient boosting -- `pymc.png`, `statsmodel.svg`: Statistical modeling - -**Special Purpose Icons:** -- `data-cleaning.png`: Data cleaning workflows -- `feature-engineering.png`: Feature engineering tasks - -## Agent Loading Process - -### For Individual Queries - -When a user makes a query like `@preprocessing_agent analyze my data`: - -1. **Check User Preferences**: System looks up user's enabled individual templates -2. **Apply Defaults**: If no preference exists, default agents are enabled -3. **Load Agent**: System loads the specific agent template and executes the query -4. **Track Usage**: Usage count is incremented for analytics - -### For Planner Workflows - -When a user makes a general query that triggers the planner: - -1. **Get Enabled Planner Templates**: System queries user's enabled planner variants -2. **Apply 10-Template Limit**: Maximum 10 templates for performance -3. **Sort by Usage**: Most-used templates get priority -4. **Create Plan**: Planner selects appropriate agents for the analysis -5. **Execute Workflow**: Selected agents execute in sequence -6. **Update Usage**: Usage statistics updated for selected agents - -### Default Agent Behavior - -```python -# Default agents enabled for new users -individual_defaults = [ - "preprocessing_agent", - "statistical_analytics_agent", - "sk_learn_agent", - "data_viz_agent" -] - -planner_defaults = [ - "planner_preprocessing_agent", - "planner_statistical_analytics_agent", - "planner_sk_learn_agent", - "planner_data_viz_agent" -] -``` - -## Usage Analytics - -### Global Usage Tracking - -The system tracks global usage statistics across all users: -- **Total usage count** per template -- **User-specific usage** for personalization -- **Last used timestamps** for sorting - -### Usage-Based Features - -- **Template Recommendations**: Popular templates shown first -- **Personalized Ordering**: User's most-used templates prioritized -- **Analytics Dashboard**: Usage patterns for administrators - -## Template Restrictions - -### User Limitations - -- **Minimum 1 Agent**: Cannot disable all templates -- **Maximum 10 for Planner**: Performance optimization -- **Premium Templates**: Some templates require premium access - -### System Limitations - -- **Active Templates Only**: Inactive templates not available -- **Variant Compatibility**: Individual/planner variants managed separately -- **Category Organization**: Templates must belong to valid categories - -## Integration with Deep Analysis - -The deep analysis system uses the template preference system: - -1. **Load User Preferences**: Gets enabled planner templates for user -2. **Create Agent Pool**: Instantiates agents from enabled templates -3. **Execute Analysis**: Uses available agents for comprehensive analysis -4. **Fallback Behavior**: Uses default agents if no preferences found - -This ensures users get personalized deep analysis based on their template preferences while maintaining system performance through the 10-template limit. \ No newline at end of file diff --git a/docs/architecture/architecture.md b/docs/architecture/architecture.md deleted file mode 100644 index 01d6dcb477c97f09f981e9c17981169b9932b78c..0000000000000000000000000000000000000000 --- a/docs/architecture/architecture.md +++ /dev/null @@ -1,427 +0,0 @@ -# Auto-Analyst Backend System Architecture - -## Overview - -Auto-Analyst is a sophisticated multi-agent AI platform designed for comprehensive data analysis. The backend system orchestrates specialized AI agents, manages user sessions, and provides a robust API for data processing and analysis workflows. - -## 🏗️ High-Level Architecture - -``` -┌─────────────────┐ ┌──────────────────┐ ┌─────────────────┐ -│ Frontend │ │ Backend │ │ Database │ -│ (Next.js) │◄──►│ (FastAPI) │◄──►│ (PostgreSQL/ │ -│ │ │ │ │ SQLite) │ -└─────────────────┘ └──────────────────┘ └─────────────────┘ - │ - ▼ - ┌──────────────────┐ - │ AI Models │ - │ (DSPy/LLMs) │ - └──────────────────┘ - │ - ▼ - ┌──────────────────┐ - │ Agent System │ - │ [Processing] │ - │ [Analytics] │ - │ [Visualization] │ - └──────────────────┘ -``` - -## 🎯 Core Components - -### 1. Application Layer (`app.py`) - -**FastAPI Application Server** -- **Role**: Main HTTP server and request router -- **Responsibilities**: - - Request/response handling - - Session-based authentication - - Route registration and middleware - - Error handling and logging - - Static file serving - - CORS configuration - -**Key Features**: -- Async/await support for high concurrency -- Automatic API documentation generation -- Request validation with Pydantic -- Session management for user tracking - -### 2. Agent System (`src/agents/`) - -**Multi-Agent Orchestra** -- **Core Agents**: Specialized AI agents for different analysis tasks -- **Deep Analysis**: Advanced multi-agent coordination system -- **Template System**: User-customizable agent configurations - -#### Agent Types - -1. **Individual Agents** (`agents.py`): - ```python - - preprocessing_agent # Data cleaning and preparation - - statistical_analytics_agent # Statistical analysis - - sk_learn_agent # Machine learning with scikit-learn - - data_viz_agent # Data visualization - - basic_qa_agent # General Q&A - ``` - -2. **Planner Agents** (Multi-agent coordination): - ```python - - planner_preprocessing_agent - - planner_statistical_analytics_agent - - planner_sk_learn_agent - - planner_data_viz_agent - ``` - -3. **Deep Analysis System** (`deep_agents.py`): - ```python - - deep_questions # Question generation - - deep_planner # Execution planning - - deep_code_synthesizer # Code combination - - deep_synthesizer # Result synthesis - - final_conclusion # Report generation - ``` - -#### Agent Architecture Pattern - -```python -class AgentSignature(dspy.Signature): - """Agent description and purpose""" - goal = dspy.InputField(desc="Analysis objective") - dataset = dspy.InputField(desc="Dataset information") - plan_instructions = dspy.InputField(desc="Execution plan") - - summary = dspy.OutputField(desc="Analysis summary") - code = dspy.OutputField(desc="Generated code") -``` - -### 3. Database Layer (`src/db/`) - -**Data Persistence and Management** - -#### Database Models (`schemas/models.py`): - -```python -# Core Models -User # User accounts and authentication -Chat # Conversation sessions -Message # Individual messages in chats -ModelUsage # AI model usage tracking - -# Template System -AgentTemplate # Agent definitions and configurations -UserTemplatePreference # User's enabled/disabled agents - -# Deep Analysis -DeepAnalysisReport # Analysis reports and results - -# Analytics -CodeExecution # Code execution tracking -UserAnalytics # User behavior analytics -``` - -#### Database Architecture: - -``` -Users (1) ──────── (Many) Chats - │ │ - │ ▼ - └─── (Many) ModelUsage ──┘ - │ - └─── (Many) UserTemplatePreference - │ - ▼ - AgentTemplate -``` - -### 4. Route Handlers (`src/routes/`) - -**RESTful API Endpoints** - -| Module | Purpose | Key Endpoints | -|--------|---------|---------------| -| `session_routes.py` | Core functionality | `/upload_excel`, `/session_info` | -| `chat_routes.py` | Chat management | `/chats`, `/messages`, `/delete_chat` | -| `code_routes.py` | Code operations | `/execute_code`, `/get_latest_code` | -| `templates_routes.py` | Agent templates | `/templates`, `/user/{id}/enabled` | -| `deep_analysis_routes.py` | Deep analysis | `/reports`, `/download_from_db` | -| `analytics_routes.py` | System analytics | `/usage`, `/feedback`, `/costs` | -| `feedback_routes.py` | User feedback | `/feedback`, `/message/{id}/feedback` | - -NOTE: Make sure to add a router prefix when calling these endpoints, such as to get dashboard, you'll use `http://localhost:8000/templates/dashboard` - -### 5. Business Logic Layer (`src/managers/`) - -**Service Layer for Complex Operations** - -#### Manager Components: - -1. **`chat_manager.py`**: - ```python - - Session management - - Message handling - - Context preservation - - Agent orchestration - ``` - -2. **`ai_manager.py`**: - ```python - - Model selection and routing - - Token tracking and cost calculation - - Error handling and retries - - Response formatting - ``` - -3. **`session_manager.py`**: - ```python - - Session lifecycle management - - Data sharing between agents - - Memory management - - Cleanup operations - ``` - -### 6. Utility Layer (`src/utils/`) - -**Shared Services and Helpers** - -- **`logger.py`**: Centralized logging system -- **`generate_report.py`**: HTML report generation -- **`model_registry.py`**: AI model configuration - -## 🔄 Data Flow Architecture - -### 1. Request Processing Flow - -``` -HTTP Request → FastAPI Router → Route Handler → Manager/Business Logic → -Database/Agent System → AI Model → Response Processing → JSON Response -``` - -### 2. Agent Execution Flow - -``` -User Query → Session Creation → Template Selection → Agent Loading → -Code Generation → Code Execution → Result Processing → Response Formatting -``` - -### 3. Deep Analysis Flow - -``` -Analysis Goal → Question Generation → Planning Phase → Agent Coordination → -Code Synthesis → Execution → Result Synthesis → Final Report Generation -``` - -### 4. Template System Flow - -``` -User Preferences → Template Loading → Agent Registration → -Capability Mapping → Execution Routing → Usage Tracking -``` - -## 🎨 Design Patterns - -### 1. **Module Pattern** -- Clear separation of concerns -- Each module has specific responsibilities -- Minimal dependencies between modules - -### 2. **Repository Pattern** -- Database access abstracted through SQLAlchemy -- Session management centralized -- Clean separation of data and business logic - -### 3. **Strategy Pattern** -- Multiple AI models supported through unified interface -- Agent selection based on user preferences -- Dynamic template loading - -### 4. **Observer Pattern** -- Usage tracking and analytics -- Event-driven model updates -- Real-time progress notifications - -## 🔧 Configuration Management - -### Environment Configuration - -```python -# Database -DATABASE_URL: str # Database connection string -POSTGRES_PASSWORD: str # PostgreSQL password (optional) - -# AI Models -ANTHROPIC_API_KEY: str # Claude API key -OPENAI_API_KEY: str # OpenAI API key - -# Authentication -ADMIN_API_KEY: str # Admin operations key (optional) - -# Deployment -PORT: int = 8000 # Server port -DEBUG: bool = False # Debug mode -``` - -### Agent Configuration (`agents_config.json`) - -```json -{ - "default_agents": [ - { - "template_name": "preprocessing_agent", - "description": "Data cleaning and preparation", - "variant_type": "both", - "is_premium": false, - "usage_count": 0, - "icon_url": "preprocessing.svg" - } - ], - "premium_templates": [...], - "remove": [...] -} -``` - -## 🔒 Security Architecture - -### Authentication & Authorization - -1. **Session-based Authentication**: - - Session IDs for user identification - - Optional API key authentication for admin endpoints - -2. **Input Validation**: - - Pydantic models for request validation - - SQL injection prevention through SQLAlchemy - - File upload restrictions and validation - -3. **Resource Protection**: - - User-specific data isolation - - Usage tracking and monitoring - - Rate limiting considerations - -### Data Security - -1. **Database Security**: - - Encrypted connections for PostgreSQL - - Parameterized queries - - Regular backup procedures - -2. **Code Execution Security**: - - Sandboxed code execution environment - - Limited library imports - - Timeout protection - -## 📊 Performance Architecture - -### Scalability Features - -1. **Async Architecture**: - - Non-blocking I/O operations - - Concurrent agent execution - - Streaming responses for long operations - -2. **Database Optimization**: - - Connection pooling - - Query optimization - - Indexed frequently accessed columns - -3. **Caching Strategy**: - - In-memory caching for templates - - Result caching for expensive operations - - Session data management - -### Performance Monitoring - -1. **Usage Analytics**: - - Request/response time tracking - - Token usage monitoring - - Error rate analysis - -2. **Resource Monitoring**: - - Database query performance - - Memory usage tracking - - Agent execution time analysis - -## 🚀 Deployment Architecture - -### Development Environment - -``` -Local Development → SQLite Database → File-based Logging → -Direct Model API Calls → Hot Reloading -``` - -### Production Environment - -``` -Load Balancer → Multiple FastAPI Instances → PostgreSQL Database → -Centralized Logging → Monitoring & Alerting -``` - -### Container Architecture - -```dockerfile -# Multi-stage build for optimization -FROM python:3.11-slim as base -# Dependencies and application setup -# Health checks and graceful shutdown -# Environment-specific configurations -``` - -## 🔄 Integration Patterns - -### External Service Integration - -1. **AI Model Providers**: - - Anthropic (Claude) - - OpenAI (GPT models) - - Unified interface through DSPy - -2. **Database Systems**: - - PostgreSQL (production) - - SQLite (development) - - Migration support through Alembic - -### Frontend Integration - -1. **REST API**: - - Standard HTTP endpoints - - JSON request/response format - - Session-based communication - -2. **Data Exchange**: - - File upload capabilities - - Real-time analysis results - - Report generation and download - -### Third-Party Integration - -1. **Python Data Science Stack (For Agentic Use only)**: - - Pandas for data manipulation - - NumPy for numerical computing - - Scikit-learn for machine learning - - Plotly for visualization - - Statsmodels for statistical analysis - -2. **Development Tools**: - - Alembic for database migrations - - SQLAlchemy for ORM - - FastAPI for web framework - - Pydantic for data validation - -## 📝 Documentation Architecture - -### API Documentation - -1. **Auto-generated Docs**: Available at `/docs` endpoint -2. **Schema Definitions**: Pydantic models with descriptions -3. **Endpoint Documentation**: Detailed parameter and response docs - -### Code Documentation - -1. **Inline Documentation**: Comprehensive docstrings -2. **Architecture Guides**: High-level system design documentation -3. **Getting Started**: Developer onboarding documentation -4. **Troubleshooting**: Common issues and solutions - -This architecture provides a robust, scalable foundation for multi-agent AI analysis while maintaining clean separation of concerns and supporting both development and production deployment scenarios. \ No newline at end of file diff --git a/docs/db_schema.md b/docs/db_schema.md new file mode 100644 index 0000000000000000000000000000000000000000..b4548b85d27f9b2902bdfb7fda4447d59cb75ef9 --- /dev/null +++ b/docs/db_schema.md @@ -0,0 +1,74 @@ +# Schema Design Explanation + +The database schema is designed to support the application's functionality by efficiently managing user data, chat sessions, messages, and model usage statistics. Each model is structured to ensure data integrity and facilitate quick access to relevant information, enabling seamless interactions within the Auto-Analyst application. + +## Database Schema + +The application uses the following database models: + +### User +The User table is used to store user information, including their username, email, and creation date. It is defined in the `init_db.py` file at [init_db.py#L17](../src/init_db.py#L17). + + +```python +class User(Base): + __tablename__ = 'users' + + user_id = Column(Integer, primary_key=True, autoincrement=True) + username = Column(String, unique=True, nullable=False) + email = Column(String, unique=True, nullable=False) + created_at = Column(DateTime, default=datetime.utcnow) +``` + +### Chat +The Chat table is used to store chat sessions, including the chat ID, user ID, title, and creation date. It is defined in the `init_db.py` file at [init_db.py#L25](../src/init_db.py#L25). + +```python +class Chat(Base): + __tablename__ = 'chats' + + chat_id = Column(Integer, primary_key=True, autoincrement=True) + user_id = Column(Integer, ForeignKey('users.user_id'), nullable=True) + title = Column(String, default='New Chat') + created_at = Column(DateTime, default=datetime.utcnow) +``` + +### Message +The Message table is used to store messages, including the message ID, chat ID, sender, content, and timestamp. It is defined in the `init_db.py` file at [init_db.py#L30](../src/init_db.py#L30). + +```python +class Message(Base): + __tablename__ = 'messages' + + message_id = Column(Integer, primary_key=True, autoincrement=True) + chat_id = Column(Integer, ForeignKey('chats.chat_id'), nullable=False) + sender = Column(String, nullable=False) # 'user' or 'ai' + content = Column(Text, nullable=False) + timestamp = Column(DateTime, default=datetime.utcnow) +``` + +### ModelUsage +The ModelUsage table is used to store model usage statistics, including the usage ID, user ID, chat ID, model name, provider, prompt tokens, completion tokens, total tokens, query size, response size, cost, timestamp, and streaming status. It is defined in the `init_db.py` file at [init_db.py#L40](../src/init_db.py#L40). + +```python +class ModelUsage(Base): + __tablename__ = 'model_usage' + + usage_id = Column(Integer, primary_key=True) + user_id = Column(Integer, ForeignKey('users.user_id'), nullable=True) + chat_id = Column(Integer, ForeignKey('chats.chat_id'), nullable=True) + model_name = Column(String(100), nullable=False) + provider = Column(String(50), nullable=False) + prompt_tokens = Column(Integer, default=0) + completion_tokens = Column(Integer, default=0) + total_tokens = Column(Integer, default=0) + query_size = Column(Integer, default=0) # Size in characters + response_size = Column(Integer, default=0) # Size in characters + cost = Column(Float, default=0.0) # Cost in USD + timestamp = Column(DateTime, default=datetime.utcnow) + is_streaming = Column(Boolean, default=False) + request_time_ms = Column(Integer, default=0) # Request processing time in milliseconds +``` + + +The models defined in this file are utilized for storing and retrieving data in the database. The database schema is managed using SQLAlchemy, which facilitates the creation and manipulation of these models. For the initialization of the database, refer to the [init_db.py](/Auto-Analyst-CS/auto-analyst-backend/src/init_db.py) file. diff --git a/docs/development/development_workflow.md b/docs/development/development_workflow.md deleted file mode 100644 index 6b3852de5458fbac8f2291faa459c9b021021ac4..0000000000000000000000000000000000000000 --- a/docs/development/development_workflow.md +++ /dev/null @@ -1,506 +0,0 @@ -# Auto-Analyst Backend Development Workflow - -## 🎯 Development Philosophy - -The Auto-Analyst backend follows modern Python development practices with emphasis on: -- **Modularity**: Clear separation of concerns across components -- **Async-First**: Non-blocking operations for scalability -- **Type Safety**: Comprehensive type hints and validation -- **Documentation**: Self-documenting code and comprehensive docs -- **Testing**: Robust testing at multiple levels -- **Performance**: Optimized for real-world usage patterns - -## 🏗️ Code Organization Principles - -### 1. **Directory Structure Standards** - -``` -src/ -├── agents/ # AI agent implementations -│ ├── agents.py # Core agent definitions -│ ├── deep_agents.py # Deep analysis system -│ └── retrievers/ # Information retrieval components -├── db/ # Database layer -│ ├── init_db.py # Database initialization -│ └── schemas/ # SQLAlchemy models -├── managers/ # Business logic layer -│ ├── chat_manager.py # Chat operations -│ ├── ai_manager.py # AI model management -│ └── session_manager.py # Session lifecycle -├── routes/ # FastAPI route handlers -│ ├── session_routes.py # Core functionality -│ ├── chat_routes.py # Chat endpoints -│ └── [feature]_routes.py # Feature-specific routes -├── utils/ # Shared utilities -│ ├── logger.py # Centralized logging -│ └── helpers.py # Common functions -└── schemas/ # Pydantic models - ├── chat_schemas.py # Chat data models - └── [feature]_schemas.py # Feature schemas -``` - -### 2. **Import Organization** - -```python -# Standard library imports -import asyncio -import json -from datetime import datetime -from typing import List, Optional, Dict, Any - -# Third-party imports -import dspy -import pandas as pd -from fastapi import APIRouter, Depends, HTTPException -from pydantic import BaseModel -from sqlalchemy.orm import Session - -# Local imports -from src.db.init_db import session_factory -from src.db.schemas.models import User, Chat -from src.utils.logger import Logger -from src.managers.chat_manager import ChatManager -``` - -## 🛠️ Development Patterns - -### 1. **Agent Development Pattern** - -```python -# 1. Define DSPy Signature -class new_analysis_agent(dspy.Signature): - """ - Comprehensive docstring explaining: - - Agent purpose and capabilities - - Input requirements and formats - - Expected output format - - Usage examples - """ - goal = dspy.InputField(desc="Clear description of analysis objective") - dataset = dspy.InputField(desc="Dataset structure and content description") - plan_instructions = dspy.InputField(desc="Execution plan from planner") - - summary = dspy.OutputField(desc="Natural language summary of analysis") - code = dspy.OutputField(desc="Executable Python code for analysis") - -# 2. Add to Agent Configuration -# In agents_config.json: -{ - "template_name": "new_analysis_agent", - "description": "Performs specialized analysis on datasets", - "variant_type": "both", # individual, planner, or both - "is_premium": false, # Will be active by default - "usage_count": 0, - "icon_url": "analysis.svg" -} - -# 3. Register in Agent System -# In agents.py, add to the appropriate loading functions -``` - -### 2. **Route Development Pattern** - -```python -# 1. Create route file: src/routes/feature_routes.py -from fastapi import APIRouter, Depends, HTTPException, Query -from pydantic import BaseModel -from typing import List, Optional -from src.db.init_db import session_factory -from src.db.schemas.models import FeatureModel -from src.utils.logger import Logger - -logger = Logger("feature_routes", see_time=True, console_log=False) -router = APIRouter(prefix="/feature", tags=["feature"]) - -# 2. Define Pydantic schemas -class FeatureCreate(BaseModel): - name: str - description: Optional[str] = None - -class FeatureResponse(BaseModel): - id: int - name: str - description: Optional[str] - created_at: datetime - -# 3. Implement endpoints with proper error handling -@router.post("/", response_model=FeatureResponse) -async def create_feature(feature: FeatureCreate): - try: - session = session_factory() - try: - new_feature = FeatureModel( - name=feature.name, - description=feature.description - ) - session.add(new_feature) - session.commit() - session.refresh(new_feature) - - return FeatureResponse( - id=new_feature.id, - name=new_feature.name, - description=new_feature.description, - created_at=new_feature.created_at - ) - - except Exception as e: - session.rollback() - logger.log_message(f"Error creating feature: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to create feature: {str(e)}") - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error in create_feature: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") - -# 4. Register in app.py -from src.routes.feature_routes import router as feature_router -app.include_router(feature_router) -``` - -### 3. **Database Model Pattern** - -```python -# In src/db/schemas/models.py -from sqlalchemy import Column, Integer, String, DateTime, Boolean, Text, ForeignKey -from sqlalchemy.orm import relationship -from sqlalchemy.ext.declarative import declarative_base -from datetime import datetime, timezone - -Base = declarative_base() - -class NewModel(Base): - __tablename__ = "new_models" - - # Primary key - id = Column(Integer, primary_key=True, autoincrement=True) - - # Required fields - name = Column(String(255), nullable=False, unique=True) - - # Optional fields - description = Column(Text, nullable=True) - is_active = Column(Boolean, default=True, nullable=False) - - # Timestamps - created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), nullable=False) - updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc), nullable=False) - - # Foreign keys - user_id = Column(Integer, ForeignKey("users.user_id"), nullable=True) - - # Relationships - user = relationship("User", back_populates="new_models") - - def __repr__(self): - return f"" - -# Update User model to include back reference -class User(Base): - # ... existing fields ... - new_models = relationship("NewModel", back_populates="user") -``` - -### 4. **Manager Pattern** - -```python -# In src/managers/feature_manager.py -from typing import List, Optional, Dict, Any -from sqlalchemy.orm import Session -from src.db.schemas.models import FeatureModel -from src.utils.logger import Logger - -logger = Logger("feature_manager", see_time=True, console_log=False) - -class FeatureManager: - """ - Manages business logic for feature operations. - Separates complex business logic from route handlers. - """ - - def __init__(self, session: Session): - self.session = session - - async def create_feature(self, name: str, description: Optional[str] = None) -> FeatureModel: - """Create a new feature with validation and business logic.""" - try: - # Validation - if not name or len(name.strip()) == 0: - raise ValueError("Feature name cannot be empty") - - # Check for duplicates - existing = self.session.query(FeatureModel).filter_by(name=name).first() - if existing: - raise ValueError(f"Feature with name '{name}' already exists") - - # Create feature - feature = FeatureModel(name=name, description=description) - self.session.add(feature) - self.session.commit() - self.session.refresh(feature) - - logger.log_message(f"Created feature: {name}", level=logging.INFO) - return feature - - except Exception as e: - self.session.rollback() - logger.log_message(f"Error creating feature: {str(e)}", level=logging.ERROR) - raise - - async def get_features(self, active_only: bool = True) -> List[FeatureModel]: - """Retrieve features with optional filtering.""" - try: - query = self.session.query(FeatureModel) - if active_only: - query = query.filter(FeatureModel.is_active == True) - - features = query.order_by(FeatureModel.created_at.desc()).all() - return features - - except Exception as e: - logger.log_message(f"Error retrieving features: {str(e)}", level=logging.ERROR) - raise -``` - -## 📋 Code Quality Standards - -### 1. **Type Hints and Documentation** - -```python -from typing import List, Optional, Dict, Any, Union -from datetime import datetime - -async def process_analysis_data( - data: pd.DataFrame, - analysis_type: str, - user_id: Optional[int] = None, - options: Dict[str, Any] = None -) -> Dict[str, Union[str, List[Any], bool]]: - """ - Process analysis data with specified parameters. - - Args: - data: Input DataFrame containing the data to analyze - analysis_type: Type of analysis to perform ("statistical", "ml", "viz") - user_id: Optional user ID for tracking and personalization - options: Additional options for analysis configuration - - Returns: - Dictionary containing: - - status: "success" or "error" - - result: Analysis results or error message - - metadata: Additional information about the analysis - - Raises: - ValueError: If analysis_type is not supported - DataError: If data format is invalid - - Example: - >>> data = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}) - >>> result = await process_analysis_data(data, "statistical") - >>> print(result["status"]) - "success" - """ - if options is None: - options = {} - - # Implementation... - return {"status": "success", "result": [], "metadata": {}} -``` - -### 2. **Error Handling Patterns** - -```python -# Comprehensive error handling with logging and user-friendly messages -async def safe_operation(data: Any) -> Dict[str, Any]: - """ - Template for safe operations with comprehensive error handling. - """ - try: - # Validation - if not data: - raise ValueError("Data cannot be empty") - - # Main operation - result = await perform_operation(data) - - # Success logging - logger.log_message("Operation completed successfully", level=logging.INFO) - return {"success": True, "data": result} - - except ValueError as e: - # Input validation errors - logger.log_message(f"Validation error: {str(e)}", level=logging.WARNING) - return {"success": False, "error": "Invalid input", "details": str(e)} - - except ConnectionError as e: - # External service errors - logger.log_message(f"Connection error: {str(e)}", level=logging.ERROR) - return {"success": False, "error": "Service unavailable", "details": "Please try again later"} - - except Exception as e: - # Unexpected errors - logger.log_message(f"Unexpected error in safe_operation: {str(e)}", level=logging.ERROR) - return {"success": False, "error": "Internal error", "details": "Please contact support"} -``` - -### 3. **Async/Await Best Practices** - -```python -import asyncio -from typing import List, Coroutine - -# Proper async function definition -async def async_agent_execution(agents: List[str], query: str) -> List[Dict[str, Any]]: - """Execute multiple agents concurrently.""" - - # Create coroutines - tasks = [ - execute_single_agent(agent, query) - for agent in agents - ] - - # Execute concurrently with error handling - results = [] - for task in asyncio.as_completed(tasks): - try: - result = await task - results.append(result) - except Exception as e: - logger.log_message(f"Agent execution failed: {e}", level=logging.ERROR) - results.append({"error": str(e)}) - - return results - -# Database operations with proper session management -async def async_database_operation(session: Session) -> Any: - """Template for async database operations.""" - try: - # Use asyncio.to_thread for CPU-bound database operations - result = await asyncio.to_thread( - lambda: session.query(Model).filter(...).all() - ) - return result - except Exception as e: - session.rollback() - raise - finally: - session.close() -``` - -## 🔧 Development Workflow - -### 1. **Feature Development Process** - -1. **Plan the Feature**: - ```bash - # Create feature branch - git checkout -b feature/new-analysis-agent - - # Document requirements - echo "## New Analysis Agent" >> docs/feature_plan.md - ``` - -2. **Implement Core Logic**: - ```bash - # Create agent signature - # Add to agents_config.json - # Implement business logic in managers/ - # Create route handlers - ``` - -3. **Add Database Changes**: - ```bash - # Modify models if needed - alembic revision --autogenerate -m "Add new analysis tables" - alembic upgrade head - ``` - -### 3. **Release Process** - -1. **Pre-release Testing**: - ```bash - # Run full test suite - pytest tests/ - - # Test database migrations - alembic upgrade head - - # Test with sample data - python scripts/test_with_sample_data.py - ``` - -2. **Documentation Updates**: - ```bash - # Update API documentation - # Update troubleshooting guide - # Update changelog - ``` - -3. **Deployment Preparation**: - ```bash - # Update requirements.txt - pip freeze > requirements.txt - - # Test container build - docker build -t auto-analyst-backend . - - ``` - -## 📊 Performance Considerations - -### 1. **Database Optimization** - -```python -# Use query optimization -from sqlalchemy.orm import joinedload - -# Bad: N+1 query problem -users = session.query(User).all() -for user in users: - print(user.chats) # Separate query for each user - -# Good: Eager loading -users = session.query(User).options(joinedload(User.chats)).all() -for user in users: - print(user.chats) # No additional queries - -# Use pagination for large datasets -def get_paginated_results(session, model, page=1, per_page=20): - offset = (page - 1) * per_page - return session.query(model).offset(offset).limit(per_page).all() -``` - - -### 2. **Async Optimization** - -```python -# Use connection pooling -from sqlalchemy.pool import QueuePool - -engine = create_engine( - DATABASE_URL, - poolclass=QueuePool, - pool_size=20, - max_overflow=30 -) - -# Batch operations -async def batch_process_agents(agents: List[str], queries: List[str]): - semaphore = asyncio.Semaphore(5) # Limit concurrent operations - - async def process_with_limit(agent, query): - async with semaphore: - return await process_agent(agent, query) - - tasks = [ - process_with_limit(agent, query) - for agent, query in zip(agents, queries) - ] - - return await asyncio.gather(*tasks, return_exceptions=True) -``` - -This development workflow guide provides a comprehensive framework for maintaining code quality, consistency, and performance in the Auto-Analyst backend system. Following these patterns ensures that new features integrate seamlessly with the existing architecture while maintaining the high standards of the codebase. \ No newline at end of file diff --git a/docs/endpoints.md b/docs/endpoints.md new file mode 100644 index 0000000000000000000000000000000000000000..0262c3f716e54734cf76cbb63aaea8c7f6c60338 --- /dev/null +++ b/docs/endpoints.md @@ -0,0 +1,10 @@ +# Auto-Analyst Backend API Documentation + +This document is a guide to the backend API endpoints utilized within the Auto-Analyst application. It encompasses a thorough breakdown of various aspects, including the handling of requests, the intricate processes of data transformations, and the structured responses that the API generates. + +The Auto-Analyst application is designed to facilitate seamless interactions and efficient data management, making it essential for users to understand the available endpoints and their functionalities. Each section of this documentation is crafted to provide clarity and insight into how the API operates, ensuring that developers and users alike can effectively leverage its capabilities. + +For more specific details regarding the various functionalities offered by the API, please refer to the following sections, which delve deeper into their respective areas: +- [Analytics Endpoints](./analytics.md): Explore the endpoints dedicated to analytics, providing insights into usage statistics and performance metrics. +- [Chat Endpoints](./chats.md): Discover the endpoints that manage chat interactions, enabling users to create, retrieve, and manage chat sessions effectively. +- [Core Endpoints](./core.md): Review the core endpoints that handle fundamental operations within the application, including data uploads and session management. diff --git a/docs/getting_started.md b/docs/getting_started.md deleted file mode 100644 index 9cd35ed87b6cf46daa7f57b34b18d14c6cd5bc43..0000000000000000000000000000000000000000 --- a/docs/getting_started.md +++ /dev/null @@ -1,273 +0,0 @@ -# Auto-Analyst Backend - Getting Started Guide - -## 🎯 Overview - -This guide will help you set up and understand the Auto-Analyst backend system. Auto-Analyst is a multi-agent AI platform that orchestrates specialized agents for comprehensive data analysis. - -## 🏗️ Core Concepts - -### 1. **Multi-Agent System** -The platform uses specialized AI agents: -- **Preprocessing Agent**: Data cleaning and preparation -- **Statistical Analytics Agent**: Statistical analysis and insights -- **Machine Learning Agent**: Scikit-learn based modeling -- **Data Visualization Agent**: Chart and plot generation - -### 2. **Template System** -- **Individual Agents**: Single-purpose agents for specific tasks -- **Planner Agents**: Multi-agent coordination for complex workflows -- **User Templates**: Customizable agent preferences -- **Default vs Premium**: Core agents available to all users - -### 3. **Session Management** -- Session-based user tracking -- Shared DataFrame context between agents -- Conversation history and code execution tracking - -### 4. **Deep Analysis System** -- Multi-step analysis workflow (questions → planning → execution → synthesis) -- Streaming progress updates -- HTML report generation - -## 🚀 Quick Start - -### 1. Installation - -```bash -# Clone and navigate to backend -cd Auto-Analyst-CS/auto-analyst-backend - -# Create virtual environment -python -m venv venv -source venv/bin/activate # Linux/Mac -# or -venv\Scripts\activate # Windows - -# Install dependencies -pip install -r requirements.txt -``` - -### 2. Environment Variables - -Create `.env` file with: - -```env -# Database -DATABASE_URL=sqlite:///./auto_analyst.db # For development -# DATABASE_URL=postgresql://user:pass@host:port/db # For production - -# AI Models -ANTHROPIC_API_KEY=your_anthropic_key_here -OPENAI_API_KEY=your_openai_key_here - -# Authentication (optional) -ADMIN_API_KEY=your_admin_key_here -``` - -### 3. Database Initialization - -```bash -# Initialize database and default agents -python -c " -from src.db.init_db import init_db -init_db() -print('✅ Database initialized successfully') -" -``` - -### 4. Start the Server - -```bash -# Development server -python app.py - -# Or with uvicorn -uvicorn app:app --reload --host 0.0.0.0 --port 8000 -``` - -### 5. Verify Setup - -Visit: `http://localhost:8000/docs` for interactive API documentation - -## 📚 Key Files to Understand - -### Core Application Files - -1. **`app.py`** - Main FastAPI application and core endpoints -2. **`src/agents/agents.py`** - Agent definitions and orchestration -3. **`src/agents/deep_agents.py`** - Deep analysis system -4. **`src/db/schemas/models.py`** - Database models -5. **`src/managers/chat_manager.py`** - Chat and session management - -### Route Files (API Endpoints) - -- **`src/routes/session_routes.py`** - File uploads, sessions, authentication -- **`src/routes/chat_routes.py`** - Chat and messaging -- **`src/routes/code_routes.py`** - Code execution and processing -- **`src/routes/templates_routes.py`** - Agent template management -- **`src/routes/deep_analysis_routes.py`** - Deep analysis reports -- **`src/routes/analytics_routes.py`** - Usage analytics and monitoring - -### Configuration Files - -- **`agents_config.json`** - Agent and template definitions -- **`requirements.txt`** - Python dependencies -- **`alembic.ini`** - Database migration configuration - -## 🔧 Development Workflow - -### 1. Adding New Agents - -```python -# 1. Define agent signature in src/agents/agents.py -class new_agent(dspy.Signature): - """Agent description""" - goal = dspy.InputField(desc="Analysis goal") - dataset = dspy.InputField(desc="Dataset info") - result = dspy.OutputField(desc="Analysis result") - -# 2. Add to agents_config.json -{ - "template_name": "new_agent", - "description": "Agent description", - "variant_type": "both", - "is_premium": false, - "usage_count": 0 -} - -# 3. Register in agent loading system -``` - -### 2. Adding New Endpoints - -```python -# 1. Create route in src/routes/feature_routes.py -from fastapi import APIRouter -router = APIRouter(prefix="/feature", tags=["feature"]) - -@router.get("/endpoint") -async def new_endpoint(): - return {"message": "Hello"} - -# 2. Register in app.py -from src.routes.feature_routes import router as feature_router -app.include_router(feature_router) -``` - -### 3. Database Changes - -```bash -# 1. Modify models in src/db/schemas/models.py -# 2. Create migration -alembic revision --autogenerate -m "description" -# 3. Apply migration -alembic upgrade head -``` - -## 🧪 Testing Your Changes - -### 1. Test API Endpoints - -```bash -# Use the interactive docs -open http://localhost:8000/docs - -# Or use curl -curl -X GET "http://localhost:8000/health" -``` - -### 2. Test Agent System - -```python -# Test individual agent -python -c " -from src.agents.agents import preprocessing_agent -import dspy -dspy.LM('anthropic/claude-sonnet-4-20250514') -agent = dspy.ChainOfThought(preprocessing_agent) -result = agent(goal='clean data', dataset='test data') -print(result) -" -``` - -### 3. Test Database Operations - -```python -# Test database -python -c " -from src.db.init_db import session_factory -from src.db.schemas.models import AgentTemplate -session = session_factory() -templates = session.query(AgentTemplate).all() -print(f'Found {len(templates)} templates') -session.close() -" -``` - -## 🔍 Common Development Tasks - -### Adding a New Feature - -1. **Plan the Feature**: Define requirements and API design -2. **Database Changes**: Add new models if needed -3. **Create Routes**: Add API endpoints in `src/routes/` -4. **Business Logic**: Add managers in `src/managers/` if complex -5. **Documentation**: Update relevant `.md` files -6. **Testing**: Test endpoints and integration - -### Debugging Issues - -1. **Check Logs**: Application logs show detailed error information -2. **Database State**: Verify data with database queries -3. **API Testing**: Use `/docs` interface for endpoint testing -4. **Agent Behavior**: Test individual agents separately - -### Performance Optimization - -1. **Database Queries**: Use SQLAlchemy query optimization -2. **Agent Execution**: Implement async patterns for agent orchestration -3. **Resource Management**: Monitor memory usage for large datasets - -## 📊 System Architecture Overview - -```mermaid -graph TD - A[Frontend Request] --> B[FastAPI Router] - B --> C[Route Handler] - C --> D[Manager Layer] - D --> E[Database Layer] - D --> F[Agent System] - F --> G[AI Models] - G --> H[Code Generation] - H --> I[Execution Environment] - I --> J[Results Processing] - J --> K[Response] - - subgraph "Agent Orchestration" - F1[Individual Agents] - F2[Planner Module] - F3[Deep Analysis] - F1 --> F2 - F2 --> F3 - end - - F --> F1 -``` - -## 📈 Template Integration - -The system uses **active user templates** for agent selection: - -### Default Agents (Always Available) -- `preprocessing_agent` (individual & planner variants) -- `statistical_analytics_agent` (individual & planner variants) -- `sk_learn_agent` (individual & planner variants) -- `data_viz_agent` (individual & planner variants) - -### Template Loading Logic -1. **Individual Agent Execution** (`@agent_name`): Loads ALL available templates -2. **Planner Execution**: Loads user's enabled templates (max 10 for performance) -3. **Deep Analysis**: Uses user's active template preferences -4. **Fallback**: Uses 4 core agents if no user preferences found - -This architecture ensures users can leverage their preferred agents while maintaining system performance and reliability. \ No newline at end of file diff --git a/docs/routes/analytics.md b/docs/routes/analytics.md new file mode 100644 index 0000000000000000000000000000000000000000..b0006ec27131999fa6fda765ed21448673b1c672 --- /dev/null +++ b/docs/routes/analytics.md @@ -0,0 +1,200 @@ + +## Analytics Routes Overview + +These routes provide comprehensive analytics functionality for the Auto-Analyst application, including dashboard summaries, user and model analytics, and cost breakdowns. + +### Authentication + +All analytics endpoints require admin authentication via an API key: + +```python +ADMIN_API_KEY = os.getenv("ADMIN_API_KEY", "default-admin-key-change-me") +``` + +The API key can be provided via: +- **Header:** `X-Admin-API-Key` +- **Query parameter:** `admin_api_key` + +--- + +### Dashboard Endpoints + +#### **GET /analytics/dashboard** +Returns a comprehensive summary of usage data for the dashboard. + +**Query Parameters:** +- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) + +**Response:** +```json +{ + "total_tokens": 123456, + "total_cost": 25.50, + "total_requests": 1000, + "total_users": 50, + "daily_usage": [ + { + "date": "2023-05-01", + "tokens": 5000, + "cost": 1.25, + "requests": 100 + } + ], + "model_usage": [ + { + "model_name": "gpt-4", + "tokens": 10000, + "cost": 10.00, + "requests": 200 + } + ], + "top_users": [ + { + "user_id": "123", + "tokens": 5000, + "cost": 5.00, + "requests": 50 + } + ], + "start_date": "2023-04-01", + "end_date": "2023-05-01" +} +``` + +**Process Flow:** +1. Parse date range from the `period` parameter. +2. Query total statistics (tokens, cost, requests, users). +3. Query daily usage broken down by date. +4. Query model usage statistics. +5. Query top users by token usage. +6. Return combined data. + +--- + +### WebSocket **/analytics/dashboard/realtime** + +WebSocket endpoint for real-time updates to dashboard data. + +**Data Flow:** +1. Client connects to WebSocket. +2. Server adds the connection to the active connections set. +3. Server broadcasts updates to all connected clients when new data arrives. +4. Connection is removed when the client disconnects. + +--- + +## User Analytics Endpoints + +### **GET /analytics/users** +Returns a list of users with their usage statistics. + +**Query Parameters:** +- `limit` (optional): Maximum number of users to return (default: `100`) +- `offset` (optional): Offset for pagination (default: `0`) + +**Response:** +```json +{ + "users": [ + { + "user_id": "123", + "tokens": 5000, + "cost": 5.00, + "requests": 50, + "first_seen": "2023-04-01T12:00:00Z", + "last_seen": "2023-05-01T12:00:00Z" + } + ], + "total": 200, + "limit": 100, + "offset": 0 +} +``` + +**Process Flow:** +1. Query user data with aggregated metrics. +2. Calculate total users for pagination. +3. Format and return data. + +--- + +### **GET /analytics/users/activity** +Returns user activity metrics over time. + +**Query Parameters:** +- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) + +**Response:** +```json +{ + "user_activity": [ + { + "date": "2023-05-01", + "activeUsers": 20, + "newUsers": 5, + "sessions": 30 + } + ] +} +``` + +**Process Flow:** +1. Parse date range from `period` parameter. +2. Get first date each user was seen (for new users count). +3. Get daily activity metrics. +4. Fill in any missing dates with zeros. +5. Return formatted data. + +--- + +## Model Analytics Endpoints + +### **GET /analytics/usage/models** +Returns model usage breakdown. + +**Query Parameters:** +- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) + +**Response:** +```json +{ + "model_usage": [ + { + "model_name": "gpt-4", + "tokens": 10000, + "cost": 10.00, + "requests": 200, + "avg_response_time": 1.5 + } + ] +} +``` + +**Process Flow:** +1. Parse date range from `period` parameter. +2. Query model usage with aggregated metrics. +3. Format and return data. + +--- + +## Cost Analytics Endpoints + +### **GET /analytics/costs/summary** +Returns a summary of costs. + +**Query Parameters:** +- `period` (optional): Time period (`7d`, `30d`, `90d`, default: `30d`) + +**Response:** +```json +{ + "totalCost": 25.50, + "totalTokens": 100000, + "totalRequests": 1000, + "avgDailyCost": 0.85, + "costPerThousandTokens": 0.255, + "daysInPeriod": 30, + "startDate": "2023-04-01", + "endDate": "2023-05-01" +} +``` diff --git a/docs/routes/chats.md b/docs/routes/chats.md new file mode 100644 index 0000000000000000000000000000000000000000..5a0dfcfec73e8e09f57e0ffa1a44e667e9578c4b --- /dev/null +++ b/docs/routes/chats.md @@ -0,0 +1,257 @@ +### Chat Routes Overview + +These routes handle chat interactions, message processing, user management, and debugging. + +--- + +### **Chat Management** + +#### **1. Create a New Chat** +**Endpoint:** `POST /chats` +**Description:** Creates a new chat session. +**Request Body:** +```json +{ + "user_id": 123, + "is_admin": false +} +``` +**Response:** +```json +{ + "chat_id": 456, + "title": "New Chat", + "created_at": "2023-05-01T12:00:00Z", + "user_id": 123 +} +``` +**Process Flow:** +1. Create a new chat for the given user. +2. Return the chat details. + +--- + +#### **2. Retrieve a Chat by ID** +**Endpoint:** `GET /chats/{chat_id}` +**Description:** Fetches a specific chat along with its messages. +**Path Parameter:** `chat_id` (ID of the chat) +**Query Parameter:** `user_id` (Optional for access control) +**Response:** +```json +{ + "chat_id": 456, + "title": "New Chat", + "created_at": "2023-05-01T12:00:00Z", + "user_id": 123, + "messages": [ + { + "message_id": 789, + "chat_id": 456, + "content": "Hello, how can I help?", + "sender": "ai", + "timestamp": "2023-05-01T12:01:00Z" + } + ] +} +``` +**Process Flow:** +1. Retrieve chat details. +2. Validate user access if `user_id` is provided. +3. Fetch all messages in the chat. +4. Return the chat with its messages. + +--- + +#### **3. List Recent Chats** +**Endpoint:** `GET /chats` +**Description:** Retrieves a list of recent chats, optionally filtered by user ID. +**Query Parameters:** +- `user_id` (Optional for filtering by user) +- `limit` (Maximum number of chats, default: 10) +- `offset` (For pagination, default: 0) +**Response:** +```json +[ + { + "chat_id": 456, + "title": "New Chat", + "created_at": "2023-05-01T12:00:00Z", + "user_id": 123 + } +] +``` +**Process Flow:** +1. Fetch recent chats. +2. Apply filters and pagination. +3. Return the list of chats. + +--- + +#### **4. Update a Chat** +**Endpoint:** `PUT /chats/{chat_id}` +**Description:** Updates a chat’s title or user ID. +**Path Parameter:** `chat_id` (ID of the chat to update) +**Request Body:** +```json +{ + "title": "Updated Chat Title", + "user_id": 123 +} +``` +**Response:** +```json +{ + "chat_id": 456, + "title": "Updated Chat Title", + "created_at": "2023-05-01T12:00:00Z", + "user_id": 123 +} +``` +**Process Flow:** +1. Update the chat’s title or user ID. +2. Return the updated details. + +--- + +#### **5. Delete a Chat** +**Endpoint:** `DELETE /chats/{chat_id}` +**Description:** Deletes a chat and all its messages. +**Path Parameter:** `chat_id` (ID of the chat to delete) +**Query Parameter:** `user_id` (Optional for access control) +**Response:** +```json +{ + "message": "Chat 456 deleted successfully" +} +``` +**Process Flow:** +1. Validate if the chat exists and if the user has access. +2. Delete the chat and associated messages. +3. Return a success message. + +--- + +#### **6. Search Chats** +**Endpoint:** `GET /chats/search/` +**Description:** Searches chats based on a query string. +**Query Parameters:** +- `query` (Search term) +- `user_id` (Optional filter) +- `limit` (Max results) +**Response:** +```json +[ + { + "chat_id": 456, + "title": "Chat about machine learning", + "created_at": "2023-05-01T12:00:00Z", + "user_id": 123 + } +] +``` +**Process Flow:** +1. Search chat titles and messages for the query. +2. Filter by `user_id` if provided. +3. Apply a result limit. +4. Return the matching chats. + +--- + +#### **7. Cleanup Empty Chats** +**Endpoint:** `POST /chats/cleanup-empty` +**Description:** Deletes empty chats for a user. +**Request Body:** +```json +{ + "user_id": 123, + "is_admin": false +} +``` +**Response:** +```json +{ + "message": "Deleted 5 empty chats" +} +``` +**Process Flow:** +1. Identify chats with no messages. +2. Delete those chats. +3. Return the count of deleted chats. + +--- + +### **Message Management** + + + +#### **1. Send a Message and Get AI Response** +**Endpoint:** `POST /chats/{chat_id}/message` +**Description:** Sends a message and receives an AI-generated response. +**Path Parameter:** `chat_id` (ID of the chat) +**Request Body:** +**Process Flow:** +1. Verify chat existence and user access. +2. Store the user message. +3. Generate an AI response via `ai_manager`. +4. Track model usage metrics. +5. Store the AI response. +6. Update the chat title if it's the first or second message. +7. Return both messages. + +--- + +### **User Management** + +#### **1. Create or Retrieve a User** +**Endpoint:** `POST /chats/users` +**Description:** Creates a new user or retrieves an existing one based on email. +**Request Body:** +```json +{ + "username": "john_doe", + "email": "john@example.com" +} +``` +**Response:** +```json +{ + "user_id": 123, + "username": "john_doe", + "email": "john@example.com", + "created_at": "2023-05-01T12:00:00Z" +} +``` +**Process Flow:** +1. Check if a user with the email exists. +2. Create a new user if not found. +3. Return the user details. + +--- + +### **Debugging** + +#### **1. Test Model Usage Tracking** +**Endpoint:** `POST /chats/debug/test-model-usage` +**Query Parameters:** +- `model_name`: Model to test +- `user_id`: Optional +**Response:** +```json +{ + "success": true, + "message": "Model usage tracking test completed", + "response": "This is a test response", + "usage_recorded": { + "usage_id": 123, + "model_name": "gpt-4", + "tokens": 50, + "cost": 0.0005, + "timestamp": "2023-05-01T12:00:00Z" + } +} +``` +**Process Flow:** +1. Generate a test prompt. +2. Call AI manager. +3. Fetch usage data. +4. Return test results. \ No newline at end of file diff --git a/docs/routes/core.md b/docs/routes/core.md new file mode 100644 index 0000000000000000000000000000000000000000..f4d107b49fccaa47ed08c0519d4be55d56098fbc --- /dev/null +++ b/docs/routes/core.md @@ -0,0 +1,242 @@ + +# **Auto-Analyst API Documentation** + +The core application routes are designed to manage the data and AI analysis capabilities of the Auto-Analyst application. + +## **1. Core Application Routes** +### **Data Management** +- **POST /upload_dataframe** + **Uploads a dataset for analysis.** + **Request:** + - `file`: CSV file + - `name`: Dataset name + - `description`: Dataset description + - `session_id`: Session ID + **Response:** + ```json + { "message": "Dataframe uploaded successfully", "session_id": "abc123" } + ``` + **Process Flow:** + - Read CSV file + - Create dataset description + - Update session with dataset + - Return success message + +- **GET /api/default-dataset** + **Gets the default dataset.** + **Query Parameters:** `session_id` + **Response:** + ```json + { + "headers": ["column1", "column2", ...], + "rows": [[val1, val2, ...], ...], + "name": "Housing Dataset", + "description": "A comprehensive dataset containing housing information..." + } + ``` + **Process Flow:** + - Reset session to use default dataset + - Format dataset preview + - Return formatted data + +- **POST /reset-session** + **Resets session to default dataset.** + **Query Parameters:** `session_id`, `name` (optional), `description` (optional) + **Response:** + ```json + { + "message": "Session reset to default dataset", + "session_id": "abc123", + "dataset": "Housing.csv" + } + ``` + **Process Flow:** + - Reset session + - Update dataset description (if provided) + - Return success message + +--- + +### **2. AI Analysis** +- **POST /chat/{agent_name}** + **Processes a query using a specific AI agent.** + **Path Parameters:** `agent_name` + **Request Body:** + ```json + { "query": "Analyze the relationship between price and size" } + ``` + **Query Parameters:** `session_id`, `user_id` (optional), `chat_id` (optional) + **Response:** + ```json + { + "agent_name": "data_viz_agent", + "query": "Analyze the relationship between price and size", + "response": "# Analysis\n\nThere appears to be a strong positive correlation...", + "session_id": "abc123" + } + ``` + **Process Flow:** + - Get session state + - Validate dataset and agent + - Execute agent query + - Format and return response + +- **POST /chat** + **Processes a query using multiple AI agents.** + **Request Body:** + ```json + { "query": "Analyze the housing data" } + ``` + **Response:** *Streaming JSON objects:* + ```json + {"agent": "data_viz_agent", "content": "# Visualization\n\n...", "status": "success"} + {"agent": "statistical_analytics_agent", "content": "# Statistical Analysis\n\n...", "status": "success"} + ``` + **Process Flow:** + - Get session state + - Validate dataset + - Generate AI analysis plan + - Execute with multiple agents and stream responses + +- **POST /execute_code** + **Executes Python code and returns results.** + **Request Body:** + ```json + { "code": "import pandas as pd\nimport matplotlib.pyplot as plt\n..." } + ``` + **Response:** + ```json + { + "output": "Execution successful", + "plotly_outputs": ["```plotly\n{\"data\": [...], \"layout\": {...}}\n```"] + } + ``` + **Process Flow:** + - Validate dataset + - Execute code + - Capture output and plots + - Return results + +--- + +### **3. Model Settings** +- **GET /api/model-settings** + **Fetches current model settings.** + **Response:** + ```json + { + "provider": "openai", + "model": "gpt-4o-mini", + "hasCustomKey": true, + "temperature": 1.0, + "maxTokens": 6000 + } + ``` +- **POST /settings/model** + **Updates model settings.** + **Request Body:** + ```json + { + "provider": "openai", + "model": "gpt-4", + "api_key": "sk-...", + "temperature": 0.7, + "max_tokens": 8000 + } + ``` + **Response:** + ```json + { "message": "Model settings updated successfully" } + ``` + **Process Flow:** + - Update model settings + - Test configuration + - Return success or error + +- **GET /agents** + **Lists available AI agents.** + **Response:** + ```json + { + "available_agents": ["data_viz_agent", "sk_learn_agent", "statistical_analytics_agent", "preprocessing_agent"], + "description": "List of available specialized agents that can be called using @agent_name" + } + ``` + +--- + +### **4. Authentication & Session Management** +- **Session ID Sources:** + - Query parameter: `session_id` + - Header: `X-Session-ID` + - Auto-generated if not provided +- **Session State Includes:** + - Current dataset + - AI system instance + - Model configuration + +#### **Admin Authentication** +- **API Key Sources:** + - Header: `X-Admin-API-Key` + - Query Parameter: `admin_api_key` +- **Validation:** + - Checked against `ADMIN_API_KEY` environment variable + +--- + +### **5. AI Agents Integration** +- **Available Agents:** + - `data_viz_agent`: Creates data visualizations (Plotly) + - `sk_learn_agent`: ML analysis (Scikit-learn) + - `statistical_analytics_agent`: Statistical analysis (StatsModels) + - `preprocessing_agent`: Data transformation + +**Integration Flow:** +- Agents are registered in `AVAILABLE_AGENTS` +- Queries are dispatched based on content +- Responses are streamed in Markdown + +--- + +### **6. Real-time Updates via WebSockets** +- **Endpoints:** + - `/analytics/dashboard/realtime` + - `/analytics/realtime` +- **Connections stored in:** + - `active_dashboard_connections` + - `active_user_connections` +- **Broadcast Function:** + ```python + async def broadcast_dashboard_update(update_data: Dict[str, Any]): + for connection in active_dashboard_connections.copy(): + try: + await connection.send_text(json.dumps(update_data)) + except Exception: + active_dashboard_connections.remove(connection) + ``` + +--- + +### **7. Error Handling** +- **Try-Except Blocks:** + ```python + try: + return result + except ValueError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error: {str(e)}") + raise HTTPException(status_code=500, detail=f"Failed: {str(e)}") + ``` +- **HTTP Exception Types:** + - `400`: Bad Request + - `401`: Unauthorized + - `403`: Forbidden + - `404`: Not Found + - `500`: Internal Server Error + +- **Logging Setup:** + ```python + logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + logger = logging.getLogger(__name__) + ``` diff --git a/docs/schema.md b/docs/schema.md new file mode 100644 index 0000000000000000000000000000000000000000..0083de719d9be89af5e58cc74960dd438bb0e0e2 --- /dev/null +++ b/docs/schema.md @@ -0,0 +1,48 @@ +# Entity-Relationship Diagram + +```mermaid +erDiagram + User ||--o{ Chat : has + User ||--o{ ModelUsage : has + Chat ||--o{ Message : has + Chat ||--o{ ModelUsage : linked_to + + User { + Integer user_id PK + String username "Unique, Not Null" + String email "Unique, Not Null" + DateTime created_at "Default=datetime.utcnow" + } + + Chat { + Integer chat_id PK + Integer user_id FK + String title "Default='New Chat'" + DateTime created_at "Default=datetime.utcnow" + } + + Message { + Integer message_id PK + Integer chat_id FK + String sender "Not Null" + Text content "Not Null" + DateTime timestamp "Default=datetime.utcnow" + } + + ModelUsage { + Integer usage_id PK + Integer user_id FK + Integer chat_id FK + String model_name "Not Null" + String provider "Not Null" + Integer prompt_tokens "Default=0" + Integer completion_tokens "Default=0" + Integer total_tokens "Default=0" + Integer query_size "Default=0" + Integer response_size "Default=0" + Float cost "Default=0.0" + DateTime timestamp "Default=datetime.utcnow" + Boolean is_streaming "Default=False" + Integer request_time_ms "Default=0" + } +``` \ No newline at end of file diff --git a/docs/system/shared_dataframe.md b/docs/shared_dataframe.md similarity index 100% rename from docs/system/shared_dataframe.md rename to docs/shared_dataframe.md diff --git a/docs/system/database-schema.md b/docs/system/database-schema.md deleted file mode 100644 index 8a328d7516f77178a416c180d3246f5e91bee29f..0000000000000000000000000000000000000000 --- a/docs/system/database-schema.md +++ /dev/null @@ -1,289 +0,0 @@ -# Auto-Analyst Database Schema Documentation - -## 📋 Overview - -The Auto-Analyst backend uses a relational database schema designed for scalability and data integrity. The schema supports both **SQLite** (development) and **PostgreSQL** (production) databases through SQLAlchemy ORM. - -### **Database Features** -- **User Management** - Authentication and user data -- **Chat System** - Conversation sessions and message history -- **AI Model Tracking** - Usage analytics and cost monitoring -- **Code Execution** - Code generation and execution tracking -- **Agent Templates** - Customizable AI agent configurations -- **Deep Analysis** - Multi-step analysis reports and results -- **User Feedback** - Rating and feedback system - ---- - -## 🗄️ Database Tables - -### **1. Users Table (`users`)** - -**Purpose**: Core user authentication and profile management - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `user_id` | `INTEGER` | PRIMARY KEY, AUTO INCREMENT | Unique user identifier | -| `username` | `STRING` | UNIQUE, NOT NULL | User's display name | -| `email` | `STRING` | UNIQUE, NOT NULL | User's email address | -| `created_at` | `DATETIME` | DEFAULT: UTC NOW | Account creation timestamp | - -**Relationships:** -- **One-to-Many**: `chats` (User → Chat sessions) -- **One-to-Many**: `usage_records` (User → Model usage tracking) -- **One-to-Many**: `deep_analysis_reports` (User → Analysis reports) -- **One-to-Many**: `template_preferences` (User → Agent preferences) - ---- - -### **2. Chats Table (`chats`)** - -**Purpose**: Conversation sessions and chat organization - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `chat_id` | `INTEGER` | PRIMARY KEY, AUTO INCREMENT | Unique chat session identifier | -| `user_id` | `INTEGER` | FOREIGN KEY → `users.user_id`, CASCADE DELETE | Chat owner (nullable for anonymous) | -| `title` | `STRING` | DEFAULT: 'New Chat' | Human-readable chat title | -| `created_at` | `DATETIME` | DEFAULT: UTC NOW | Chat creation timestamp | - -**Relationships:** -- **Many-to-One**: `user` (Chat → User) -- **One-to-Many**: `messages` (Chat → Messages) -- **One-to-Many**: `usage_records` (Chat → Model usage) - ---- - -### **3. Messages Table (`messages`)** - -**Purpose**: Individual messages within chat conversations - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `message_id` | `INTEGER` | PRIMARY KEY, AUTO INCREMENT | Unique message identifier | -| `chat_id` | `INTEGER` | FOREIGN KEY → `chats.chat_id`, CASCADE DELETE | Parent chat session | -| `sender` | `STRING` | NOT NULL | Message sender: 'user' or 'ai' | -| `content` | `TEXT` | NOT NULL | Message content (text/markdown) | -| `timestamp` | `DATETIME` | DEFAULT: UTC NOW | Message creation time | - -**Relationships:** -- **Many-to-One**: `chat` (Message → Chat) -- **One-to-One**: `feedback` (Message → Feedback) - ---- - -### **4. Model Usage Table (`model_usage`)** - -**Purpose**: AI model usage tracking for analytics and billing - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `usage_id` | `INTEGER` | PRIMARY KEY | Unique usage record identifier | -| `user_id` | `INTEGER` | FOREIGN KEY → `users.user_id`, SET NULL | User who triggered the usage | -| `chat_id` | `INTEGER` | FOREIGN KEY → `chats.chat_id`, SET NULL | Associated chat session | -| `model_name` | `STRING(100)` | NOT NULL | AI model used (e.g., 'gpt-4o-mini') | -| `provider` | `STRING(50)` | NOT NULL | Model provider ('openai', 'anthropic', etc.) | -| `prompt_tokens` | `INTEGER` | DEFAULT: 0 | Input tokens consumed | -| `completion_tokens` | `INTEGER` | DEFAULT: 0 | Output tokens generated | -| `total_tokens` | `INTEGER` | DEFAULT: 0 | Total tokens (input + output) | -| `query_size` | `INTEGER` | DEFAULT: 0 | Query size in characters | -| `response_size` | `INTEGER` | DEFAULT: 0 | Response size in characters | -| `cost` | `FLOAT` | DEFAULT: 0.0 | Cost in USD for this usage | -| `timestamp` | `DATETIME` | DEFAULT: UTC NOW | Usage timestamp | -| `is_streaming` | `BOOLEAN` | DEFAULT: FALSE | Whether response was streamed | -| `request_time_ms` | `INTEGER` | DEFAULT: 0 | Request processing time (milliseconds) | - -**Relationships:** -- **Many-to-One**: `user` (Usage → User) -- **Many-to-One**: `chat` (Usage → Chat) - ---- - -### **5. Code Executions Table (`code_executions`)** - -**Purpose**: Track code generation and execution attempts - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `execution_id` | `INTEGER` | PRIMARY KEY, AUTO INCREMENT | Unique execution identifier | -| `message_id` | `INTEGER` | FOREIGN KEY → `messages.message_id`, CASCADE DELETE | Associated message | -| `chat_id` | `INTEGER` | FOREIGN KEY → `chats.chat_id`, CASCADE DELETE | Parent chat session | -| `user_id` | `INTEGER` | FOREIGN KEY → `users.user_id`, SET NULL | User who triggered execution | -| `initial_code` | `TEXT` | NULLABLE | First version of generated code | -| `latest_code` | `TEXT` | NULLABLE | Most recent code version | -| `is_successful` | `BOOLEAN` | DEFAULT: FALSE | Whether execution succeeded | -| `output` | `TEXT` | NULLABLE | Execution output (including errors) | -| `model_provider` | `STRING(50)` | NULLABLE | AI model provider used | -| `model_name` | `STRING(100)` | NULLABLE | AI model name used | -| `failed_agents` | `TEXT` | NULLABLE | JSON list of failed agent names | -| `error_messages` | `TEXT` | NULLABLE | JSON map of error messages by agent | -| `created_at` | `DATETIME` | DEFAULT: UTC NOW | Execution creation time | -| `updated_at` | `DATETIME` | DEFAULT: UTC NOW, ON UPDATE | Last update timestamp | - ---- - -### **6. Message Feedback Table (`message_feedback`)** - -**Purpose**: User feedback and model settings for messages - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `feedback_id` | `INTEGER` | PRIMARY KEY, AUTO INCREMENT | Unique feedback identifier | -| `message_id` | `INTEGER` | FOREIGN KEY → `messages.message_id`, CASCADE DELETE | Associated message | -| `rating` | `INTEGER` | NULLABLE | Star rating (1-5 scale) | -| `model_name` | `STRING(100)` | NULLABLE | Model used for this message | -| `model_provider` | `STRING(50)` | NULLABLE | Model provider used | -| `temperature` | `FLOAT` | NULLABLE | Temperature setting used | -| `max_tokens` | `INTEGER` | NULLABLE | Max tokens setting used | -| `created_at` | `DATETIME` | DEFAULT: UTC NOW | Feedback creation time | -| `updated_at` | `DATETIME` | DEFAULT: UTC NOW, ON UPDATE | Last update timestamp | - -**Relationships:** -- **One-to-One**: `message` (Feedback ↔ Message) - ---- - -### **7. Deep Analysis Reports Table (`deep_analysis_reports`)** - -**Purpose**: Store comprehensive multi-agent analysis reports - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `report_id` | `INTEGER` | PRIMARY KEY, AUTO INCREMENT | Unique report identifier | -| `report_uuid` | `STRING(100)` | UNIQUE, NOT NULL | Frontend-generated UUID | -| `user_id` | `INTEGER` | FOREIGN KEY → `users.user_id`, CASCADE DELETE | Report owner | -| `goal` | `TEXT` | NOT NULL | Analysis objective/question | -| `status` | `STRING(20)` | NOT NULL, DEFAULT: 'pending' | Status: 'pending', 'running', 'completed', 'failed' | -| `start_time` | `DATETIME` | DEFAULT: UTC NOW | Analysis start time | -| `end_time` | `DATETIME` | NULLABLE | Analysis completion time | -| `duration_seconds` | `INTEGER` | NULLABLE | Total analysis duration | -| `deep_questions` | `TEXT` | NULLABLE | Generated analytical questions | -| `deep_plan` | `TEXT` | NULLABLE | Analysis execution plan | -| `summaries` | `JSON` | NULLABLE | Array of analysis summaries | -| `analysis_code` | `TEXT` | NULLABLE | Generated Python code | -| `plotly_figures` | `JSON` | NULLABLE | Array of Plotly figure data | -| `synthesis` | `JSON` | NULLABLE | Array of synthesis insights | -| `final_conclusion` | `TEXT` | NULLABLE | Final analysis conclusion | -| `html_report` | `TEXT` | NULLABLE | Complete HTML report | -| `progress_percentage` | `INTEGER` | DEFAULT: 0 | Progress percentage (0-100) | -| `total_tokens_used` | `INTEGER` | DEFAULT: 0 | Total tokens consumed | -| `estimated_cost` | `FLOAT` | DEFAULT: 0.0 | Estimated cost in USD | -| `credits_consumed` | `INTEGER` | DEFAULT: 0 | Credits deducted for analysis | -| `created_at` | `DATETIME` | DEFAULT: UTC NOW | Report creation time | -| `updated_at` | `DATETIME` | DEFAULT: UTC NOW, ON UPDATE | Last update timestamp | - -**Relationships:** -- **Many-to-One**: `user` (Report → User) - ---- - -### **8. Agent Templates Table (`agent_templates`)** - -**Purpose**: Store predefined AI agent configurations - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `template_id` | `INTEGER` | PRIMARY KEY, AUTO INCREMENT | Unique template identifier | -| `template_name` | `STRING(100)` | UNIQUE, NOT NULL | Internal template name | -| `display_name` | `STRING(200)` | NULLABLE | User-friendly display name | -| `description` | `TEXT` | NOT NULL | Template description | -| `prompt_template` | `TEXT` | NOT NULL | Agent behavior instructions | -| `icon_url` | `STRING(500)` | NULLABLE | Template icon URL | -| `category` | `STRING(50)` | NULLABLE | Template category | -| `is_premium_only` | `BOOLEAN` | DEFAULT: FALSE | Requires premium subscription | -| `variant_type` | `STRING(20)` | DEFAULT: 'individual' | 'planner', 'individual', or 'both' | -| `is_active` | `BOOLEAN` | DEFAULT: TRUE | Template is active/available | -| `created_at` | `DATETIME` | DEFAULT: UTC NOW | Template creation time | - -**Relationships:** -- **One-to-Many**: `user_preferences` (Template → User preferences) - ---- - -### **9. User Template Preferences Table (`user_template_preferences`)** - -**Purpose**: Track user preferences and usage for agent templates - -| Column | Type | Constraints | Description | -|--------|------|-------------|-------------| -| `preference_id` | `INTEGER` | PRIMARY KEY, AUTO INCREMENT | Unique preference identifier | -| `user_id` | `INTEGER` | FOREIGN KEY → `users.user_id`, CASCADE DELETE | User who owns preference | -| `template_id` | `INTEGER` | FOREIGN KEY → `agent_templates.template_id`, CASCADE DELETE | Associated template | -| `is_enabled` | `BOOLEAN` | DEFAULT: TRUE | Whether user has template enabled | -| `usage_count` | `INTEGER` | DEFAULT: 0 | Number of times user used template | -| `last_used_at` | `DATETIME` | NULLABLE | Last time user used template | -| `created_at` | `DATETIME` | DEFAULT: UTC NOW | Preference creation time | - -**Relationships:** -- **Many-to-One**: `user` (Preference → User) -- **Many-to-One**: `template` (Preference → Template) - -**Constraints:** -- **Unique**: `(user_id, template_id)` - One preference per user per template - ---- - -## 🔗 Entity Relationship Diagram - -``` -Users (1) ──────────── (Many) Chats - │ │ - │ ├── (Many) Messages - │ │ │ - │ │ └── (1) MessageFeedback - │ │ - │ └── (Many) CodeExecutions - │ - ├── (Many) ModelUsage - │ - ├── (Many) DeepAnalysisReports - │ - └── (Many) UserTemplatePreferences - │ - └── (Many) AgentTemplates -``` - ---- - -## 📊 Database Performance - -### **Optimized Indexes** - -```sql --- High-performance queries -CREATE INDEX idx_messages_chat_timestamp ON messages(chat_id, timestamp DESC); -CREATE INDEX idx_model_usage_user_time ON model_usage(user_id, timestamp DESC); -CREATE INDEX idx_model_usage_model_time ON model_usage(model_name, timestamp DESC); -CREATE INDEX idx_reports_user_time ON deep_analysis_reports(user_id, created_at DESC); -``` - -### **Cascade Deletion Rules** - -| Parent → Child | Rule | Description | -|----------------|------|-------------| -| `users` → `chats` | CASCADE | Delete all user chats when user deleted | -| `chats` → `messages` | CASCADE | Delete all chat messages when chat deleted | -| `messages` → `feedback` | CASCADE | Delete feedback when message deleted | -| `users` → `model_usage` | SET NULL | Keep usage records for analytics | - ---- - -## 🛡️ Security & Maintenance - -### **Data Protection** -- User data isolated by `user_id` -- Sensitive fields require encryption in production -- Automatic cleanup of anonymous data after 90 days - -### **Regular Maintenance** -```sql --- Clean old anonymous chats -DELETE FROM chats WHERE user_id IS NULL AND created_at < DATE_SUB(NOW(), INTERVAL 90 DAY); - --- Update statistics for query optimization -ANALYZE users, chats, messages, model_usage; -``` - ---- - -This schema supports the full Auto-Analyst application with optimized performance, data integrity, and scalability for both development and production environments. \ No newline at end of file diff --git a/docs/troubleshooting/troubleshooting.md b/docs/troubleshooting/troubleshooting.md deleted file mode 100644 index 620c5ebacef072eecc66988bb201e007a2dd5eb8..0000000000000000000000000000000000000000 --- a/docs/troubleshooting/troubleshooting.md +++ /dev/null @@ -1,537 +0,0 @@ -# Auto-Analyst Backend Troubleshooting Guide - -## 🚨 Common Startup Issues - -### 1. **Database Connection Problems** - -#### Problem: Database connection failed -``` -❌ sqlalchemy.exc.OperationalError: (sqlite3.OperationalError) no such table: users -``` - -**Solutions:** -1. **Initialize Database**: - ```bash - python -c " - from src.db.init_db import init_db - init_db() - print('✅ Database initialized') - " - ``` - -2. **Check Database File Permissions**: - ```bash - # For SQLite - ls -la auto_analyst.db - chmod 666 auto_analyst.db # If needed - ``` - -3. **Verify DATABASE_URL**: - ```bash - # Check .env file - cat .env | grep DATABASE_URL - - # For PostgreSQL (production) - DATABASE_URL=postgresql://user:password@host:port/database - - # For SQLite (development) - DATABASE_URL=sqlite:///./auto_analyst.db - ``` - -#### Problem: PostgreSQL connection issues -``` -❌ psycopg2.OperationalError: FATAL: database "auto_analyst" does not exist -``` - -**Solutions:** -1. **Create Database**: - ```sql - -- Connect to PostgreSQL - psql -h localhost -U postgres - CREATE DATABASE auto_analyst; - \q - ``` - -2. **Update Connection String**: - ```env - DATABASE_URL=postgresql://username:password@localhost:5432/auto_analyst - ``` - -### 2. **Agent Template Loading Issues** - -#### Problem: No agents found -``` -❌ RuntimeError: No agents loaded for user. Cannot proceed with analysis. -``` - -**Solutions:** -1. **Initialize Default Agents**: - ```python - python -m scripts.populate_agent_templates - print('✅ Default agents initialized') - " - ``` - -2. **Check Agent Templates in Database**: - ```python - python -c " - from src.db.init_db import session_factory - from src.db.schemas.models import AgentTemplate - session = session_factory() - templates = session.query(AgentTemplate).all() - print(f'Found {len(templates)} templates:') - for t in templates: - print(f' - {t.template_name}: {t.is_active}') - session.close() - " - ``` - -3. **Populate Templates from Config**: - ```bash - python scripts/populate_agent_templates.py - ``` - -### 3. **API Key Issues** - -#### Problem: Missing API keys -``` -❌ AuthenticationError: Invalid API key provided -``` - -**Solutions:** -1. **Check Environment Variables**: - ```bash - # Verify API keys are set - echo $ANTHROPIC_API_KEY - echo $OPENAI_API_KEY - - # Or check .env file - cat .env | grep API_KEY - ``` - -2. **Add Missing Keys**: - ```env - # Add to .env file - ANTHROPIC_API_KEY=sk-ant-api03-... - OPENAI_API_KEY=sk-... - ADMIN_API_KEY=your_admin_key_here - ``` - -3. **Test API Key Validity**: - ```python - python -c " - import os - from anthropic import Anthropic - client = Anthropic(api_key=os.getenv('ANTHROPIC_API_KEY')) - try: - # Test call - response = client.messages.create( - model='claude-3-sonnet-20241022', - max_tokens=10, - messages=[{'role': 'user', 'content': 'Hello'}] - ) - print('✅ Anthropic API key valid') - except Exception as e: - print(f'❌ Anthropic API key invalid: {e}') - " - ``` - -## 🤖 Agent System Issues - -### 1. **Agent Not Found Errors** - -#### Problem: Specific agent not available -``` -❌ KeyError: 'custom_agent' not found in loaded agents -``` - -**Solutions:** -1. **Check Available Agents**: - ```python - python -c " - from src.agents.agents import load_user_enabled_templates_from_db - from src.db.init_db import session_factory - session = session_factory() - agents = load_user_enabled_templates_from_db('test_user', session) - print('Available agents:', list(agents.keys())) - session.close() - " - ``` - -2. **Verify Agent Template Exists**: - ```python - python -c " - from src.db.init_db import session_factory - from src.db.schemas.models import AgentTemplate - session = session_factory() - agent = session.query(AgentTemplate).filter_by(template_name='custom_agent').first() - if agent: - print(f'Agent found: {agent.display_name}, Active: {agent.is_active}') - else: - print('Agent not found in database') - session.close() - " - ``` - -3. **Add Missing Agent Template**: - ```python - # Add to agents_config.json or use database insertion - python scripts/populate_agent_templates.py - ``` - -### 2. **Deep Analysis Failures** - -#### Problem: Deep analysis stops unexpectedly -``` -❌ DeepAnalysisError: Agent execution failed at step 3 -``` - -**Solutions:** -1. **Check Agent Configuration**: - ```python - # Verify user has required agents enabled - python -c " - from src.agents.deep_agents import get_user_enabled_agent_names - from src.db.init_db import session_factory - session = session_factory() - agents = get_user_enabled_agent_names('test_user', session) - required = ['preprocessing_agent', 'statistical_analytics_agent', 'sk_learn_agent', 'data_viz_agent'] - print('Required agents:', required) - print('Available agents:', agents) - print('Missing:', [a for a in required if a not in agents]) - session.close() - " - ``` - -2. **Increase Timeout Settings**: - ```python - # In deep_agents.py, increase timeout values - timeout = 300 # Increase from default - ``` - -3. **Check Dataset Size**: - ```python - # Reduce dataset size for complex analysis - df_sample = df.sample(n=1000) # Use sample for testing - ``` - -## ⚡ Code Execution Problems - -### 1. **Code Execution Timeouts** - -#### Problem: Code execution takes too long -``` -❌ TimeoutError: Code execution exceeded 120 seconds -``` - -**Solutions:** -1. **Optimize Generated Code**: - - Use data sampling for large datasets - - Simplify analysis requirements - - Use sampling for large datasets - -2. **Check Resource Usage**: - ```python - import psutil - print(f"Memory usage: {psutil.virtual_memory().percent}%") - print(f"CPU usage: {psutil.cpu_percent()}%") - ``` - -3. **Increase Timeout Settings**: - ```python - # In clean_and_store_code function - future.result(timeout=600) # Increase timeout to 10 minutes - ``` - -#### Problem: Import Errors in Generated Code -``` -❌ ModuleNotFoundError: No module named 'some_library' -``` - -**Solutions:** -1. **Check Available Libraries**: - ```python - # Available in execution environment: - import pandas as pd - import numpy as np - import plotly.express as px - import plotly.graph_objects as go - import sklearn - import statsmodels.api as sm - ``` - -2. **Add Missing Dependencies**: - ```bash - pip install missing_library - ``` - -3. **Update Execution Environment**: - ```python - # In clean_and_store_code function - exec_globals.update({ - 'new_library': __import__('new_library') - }) - ``` - -### 4. **Database Issues** - -#### Problem: Migration Errors -``` -❌ alembic.util.exc.CommandError: Can't locate revision identified by 'xyz' -``` - -**Solutions:** -1. **Reset Migration History**: - ```bash - # Delete migration files (except __init__.py) - rm migrations/versions/*.py - - # Create new initial migration - alembic revision --autogenerate -m "initial migration" - alembic upgrade head - ``` - -2. **Force Migration**: - ```bash - # Mark current state as up-to-date - alembic stamp head - ``` - -3. **Recreate Database**: - ```bash - # For SQLite (development) - rm auto_analyst.db - python -c "from src.db.init_db import init_db; init_db()" - ``` - -#### Problem: Constraint Violations -``` -❌ IntegrityError: UNIQUE constraint failed -``` - -**Solutions:** -1. **Check Existing Records**: - ```python - from src.db.init_db import session_factory - from src.db.schemas.models import AgentTemplate - - session = session_factory() - templates = session.query(AgentTemplate).all() - for t in templates: - print(f"{t.template_name}: {t.template_id}") - session.close() - ``` - -2. **Clean Duplicate Data**: - ```bash - python -c " - from src.db.init_db import session_factory - from src.db.schemas.models import AgentTemplate - session = session_factory() - # Remove duplicates based on template_name - seen = set() - for template in session.query(AgentTemplate).all(): - if template.template_name in seen: - session.delete(template) - else: - seen.add(template.template_name) - session.commit() - session.close() - " - ``` - -### 5. **Authentication and Authorization Issues** - -#### Problem: Unauthorized Access -``` -❌ 401 Unauthorized: Invalid session -``` - -**Solutions:** -1. **Check Session ID**: - ```python - # Ensure session_id is provided in request - headers = {"X-Session-ID": "your_session_id"} - # Or as query parameter: ?session_id=your_session_id - ``` - -2. **Create Valid Session**: - ```bash - curl -X POST "http://localhost:8000/session_info" \ - -H "Content-Type: application/json" - ``` - -3. **Verify Admin API Key**: - ```bash - curl -X GET "http://localhost:8000/analytics/usage" \ - -H "X-API-Key: your_admin_key" - ``` - -### 6. **Performance Issues** - -#### Problem: Slow Response Times -``` -⚠️ Request taking longer than expected -``` - -**Solutions:** -1. **Enable Database Connection Pooling**: - ```python - # In init_db.py - engine = create_engine( - DATABASE_URL, - poolclass=QueuePool, - pool_size=10, - max_overflow=20 - ) - ``` - -2. **Optimize Database Queries**: - ```python - # Use eager loading for relationships - session.query(User).options(joinedload(User.chats)).all() - ``` - -3. **Add Response Caching**: - ```python - # Use local caching for expensive operations - @lru_cache(maxsize=100) - def expensive_operation(data): - return result - ``` - -#### Problem: Memory Usage High -``` -⚠️ Memory usage above 80% -``` - -**Solutions:** -1. **Optimize DataFrame Operations**: - ```python - # Use chunking for large datasets - for chunk in pd.read_csv('file.csv', chunksize=1000): - process_chunk(chunk) - ``` - -2. **Clear Unused Variables**: - ```python - # In code execution - del large_dataframe - import gc - gc.collect() - ``` - -3. **Monitor Memory Usage**: - ```python - import psutil - import logging - - memory_percent = psutil.virtual_memory().percent - if memory_percent > 80: - logging.warning(f"High memory usage: {memory_percent}%") - ``` - -## 🔧 Debugging Tools and Commands - -### Health Check Commands - -```bash -# Test basic connectivity -curl http://localhost:8000/health - -# Check database status -python -c " -from src.db.init_db import session_factory -try: - session = session_factory() - session.execute('SELECT 1') - print('✅ Database connection OK') - session.close() -except Exception as e: - print(f'❌ Database error: {e}') -" - -# Verify agent templates -python -c " -from src.db.init_db import session_factory -from src.db.schemas.models import AgentTemplate -session = session_factory() -count = session.query(AgentTemplate).count() -print(f'Agent templates in database: {count}') -session.close() -" -``` - -### Performance Monitoring - -```python -# Memory and CPU monitoring -import psutil -import time - -def monitor_system(): - while True: - cpu = psutil.cpu_percent(interval=1) - memory = psutil.virtual_memory() - print(f"CPU: {cpu}% | Memory: {memory.percent}% | Available: {memory.available // 1024 // 1024}MB") - time.sleep(5) - -# Run monitoring -monitor_system() -``` - -### Database Inspection - -```python -# Inspect database tables -from src.db.init_db import session_factory -from src.db.schemas.models import * - -session = session_factory() - -# Count records in each table -tables = [User, Chat, Message, AgentTemplate, UserTemplatePreference, DeepAnalysisReport] -for table in tables: - count = session.query(table).count() - print(f"{table.__name__}: {count} records") - -session.close() -``` - -### Log Analysis - -```bash -# View recent logs -tail -f logs/app.log - -# Search for errors -grep "ERROR" logs/app.log | tail -20 - -# Search for specific issues -grep -i "agent" logs/app.log | grep -i "error" -``` - -## 🚀 Performance Optimization Tips - -### Database Optimization - -1. **Use Indexes**: Ensure frequently queried columns have indexes -2. **Query Optimization**: Use `joinedload` for relationships -3. **Connection Pooling**: Configure appropriate pool sizes -4. **Batch Operations**: Use bulk operations for multiple records - -### Agent Performance - -1. **Async Execution**: Use async patterns for concurrent operations -2. **Result Caching**: Cache expensive computations -3. **Memory Management**: Clean up large objects after use -4. **Code Optimization**: Simplify generated code for better performance - -### System Monitoring - -1. **Resource Tracking**: Monitor CPU, memory, and disk usage -2. **Error Monitoring**: Set up alerting for critical errors -3. **Performance Metrics**: Track response times and throughput -4. **Usage Analytics**: Monitor feature usage and optimization opportunities - -This troubleshooting guide covers the most common issues you'll encounter with the Auto-Analyst backend. For additional help, check the system logs and use the debugging tools provided. \ No newline at end of file diff --git a/entrypoint_local.sh b/entrypoint_local.sh deleted file mode 100644 index 4ac7c091d0fccb1ae10e16ba194e7b1fc7f3c749..0000000000000000000000000000000000000000 --- a/entrypoint_local.sh +++ /dev/null @@ -1,175 +0,0 @@ -#!/bin/bash - -# Entrypoint script for Auto-Analyst backend -# This script safely initializes the database and starts the application -# SAFE for PostgreSQL/RDS - only modifies SQLite databases - -set -e # Exit on any error - -echo "🚀 Starting Auto-Analyst Backend..." - -# Function to run safe database initialization -init_production_database() { - echo "🔧 Running SAFE database initialization..." - - # Run the safe initialization script - python scripts/init_production_db.py - - # Don't fail if database initialization has issues - let app try to start - if [ $? -eq 0 ]; then - echo "✅ Database initialization completed successfully" - else - echo "⚠️ Database initialization had issues, but continuing..." - echo "📋 App will start but some features may not work properly" - fi -} - -# Function to verify basic app imports work -verify_app_imports() { - echo "🔍 Verifying application imports..." - python -c " -try: - from app import app - print('✅ Main application imports successful') -except Exception as e: - print(f'❌ Application import failed: {e}') - exit(1) -" || { - echo "❌ Critical application import failure - cannot start" - exit 1 -} -} - -# Function to verify database connectivity (non-failing) -verify_database_connectivity() { - echo "🔗 Testing database connectivity..." - python -c " -try: - from src.db.init_db import get_session, is_postgres_db - from src.db.schemas.models import AgentTemplate - - db_type = 'PostgreSQL/RDS' if is_postgres_db() else 'SQLite' - print(f'🗄️ Database type: {db_type}') - - session = get_session() - - # Try to query templates if table exists - try: - template_count = session.query(AgentTemplate).count() - print(f'✅ Database connected. Found {template_count} templates.') - except Exception as table_error: - print(f'⚠️ Database connected but template table issue: {table_error}') - print('📋 Template functionality may not work') - finally: - session.close() - -except Exception as e: - print(f'⚠️ Database connectivity issue: {e}') - print('📋 App will start but database features may not work') -" - # Don't exit on database connectivity issues - let app try to start -} - -# Function to populate agents and templates for development (SQLite only) -# Uses agents_config.json if available, falls back to legacy method -populate_agents_templates() { - echo "🔧 Checking if agents/templates need to be populated..." - python -c " -try: - from src.db.init_db import DATABASE_URL - from src.db.schemas.models import AgentTemplate - from src.db.init_db import session_factory - - # Check database type - if DATABASE_URL.startswith('sqlite'): - print('🔍 SQLite database detected - checking template population') - - session = session_factory() - try: - template_count = session.query(AgentTemplate).count() - - if template_count == 0: - print('📋 No templates found - populating agents and templates...') - session.close() - exit(1) # Signal that population is needed - else: - print(f'✅ Found {template_count} templates - population not needed') - session.close() - exit(0) # Signal that population is not needed - except Exception as e: - print(f'⚠️ Error checking templates: {e}') - print('📋 Will attempt to populate anyway') - session.close() - exit(1) # Signal that population is needed - else: - print('🔍 PostgreSQL/RDS detected - skipping auto-population') - exit(0) # Signal that population is not needed - -except Exception as e: - print(f'❌ Error during template check: {e}') - exit(0) # Don't fail startup, just skip population -" - - # Check if population is needed (exit code 1 means yes) - if [ $? -eq 1 ]; then - echo "🚀 Running agent/template population for SQLite..." - - # Check if agents_config.json exists (try multiple locations) - if [ -f "agents_config.json" ] || [ -f "/app/agents_config.json" ] || [ -f "../agents_config.json" ]; then - echo "📖 Found agents_config.json - validating configuration..." - - # Validate configuration first - python scripts/populate_agent_templates.py validate - validation_result=$? - - if [ $validation_result -eq 0 ]; then - echo "✅ Configuration valid - proceeding with sync" - python scripts/populate_agent_templates.py sync - else - echo "⚠️ Configuration validation failed - attempting sync anyway" - python scripts/populate_agent_templates.py sync - fi - else - echo "⚠️ agents_config.json not found - trying legacy method" - python scripts/populate_agent_templates.py - fi - - if [ $? -eq 0 ]; then - echo "✅ Agent/template population completed successfully" - else - echo "⚠️ Agent/template population had issues, but continuing..." - echo "📋 You may need to populate templates manually" - echo "💡 Tip: Ensure agents_config.json exists in the backend directory" - fi - fi -} - -# Check if we need to find agents_config.json from space root -if [ ! -f "/app/agents_config.json" ]; then - echo "⚠️ agents_config.json not found in container - checking build issues" - echo "📁 Files in /app directory:" - ls -la /app/ | head -10 -else - echo "✅ agents_config.json found in container" -fi - -# Main startup sequence -echo "🔧 Initializing production environment..." - -# Verify critical imports first -verify_app_imports - -# Initialize database safely (won't modify RDS) -init_production_database - -# Test database connectivity (non-failing) -verify_database_connectivity - -# Populate agents and templates for development (SQLite only) -populate_agents_templates - -echo "🎯 Starting FastAPI application..." -echo "🌐 Application will be available on port 7860" - -# Start the FastAPI application -exec uvicorn app:app --host 0.0.0.0 --port 7860 \ No newline at end of file diff --git a/images/Auto-Analyst Banner.png b/images/Auto-Analyst Banner.png new file mode 100644 index 0000000000000000000000000000000000000000..c12cab9ee36147b5340e542af289ad141ed32e82 --- /dev/null +++ b/images/Auto-Analyst Banner.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:30a322031c1e8eca20f202d2bda534a921c5c4556ee9ae81f0fcabdb156ab2cd +size 184040 diff --git a/images/Auto-analysts icon small.png b/images/Auto-analysts icon small.png new file mode 100644 index 0000000000000000000000000000000000000000..49fb93737f8ff71f13134d740226ee99f6c97cd8 Binary files /dev/null and b/images/Auto-analysts icon small.png differ diff --git a/images/auto-analyst logo.png b/images/auto-analyst logo.png new file mode 100644 index 0000000000000000000000000000000000000000..b4abbb442cc261285f8716cbfa507802fbaef723 Binary files /dev/null and b/images/auto-analyst logo.png differ diff --git a/requirements.txt b/requirements.txt index 736a8f77904c7079c7f06575cdc4d4dda9b844b3..0b32185fcde2fcd56bc3b460e040e443b4e1c8b0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,31 +1,32 @@ -aiofiles==24.1.0 -beautifulsoup4==4.13.4 -chardet==5.2.0 -dspy==3.2.1 -litellm==1.82.3 +groq==0.18.0 +dspy==2.5.43 +dspy-ai==2.5.43 email_validator==2.2.0 -fastapi==0.135.1 -fastapi-cli==0.0.7 +fastapi==0.111.1 FastAPI-SQLAlchemy==0.2.1 -fastapi-sso==0.16.0 -groq==0.18.0 -gunicorn==23.0.0 -huggingface-hub==0.30.2 +gunicorn==22.0.0 joblib==1.4.2 -Markdown==3.7 +litellm==1.53.7 +llama-index==0.12.14 +llama-index-agent-openai==0.4.2 +llama-index-embeddings-openai==0.3.1 +llama-index-indices-managed-llama-cloud==0.6.4 +llama-index-llms-openai==0.3.14 +llama-index-multi-modal-llms-openai==0.4.2 +llama-index-program-openai==0.3.1 +llama-index-question-gen-openai==0.3.0 matplotlib==3.10.0 -matplotlib-inline==0.1.7 +multiprocess==0.70.16 numpy==2.2.2 openpyxl==3.1.2 xlrd==2.0.1 -openai==2.28.0 +openai==1.60.1 pandas==2.2.3 -polars==1.31.0 pillow==11.1.0 plotly==5.24.1 -psycopg2==2.9.10 -python-dateutil==2.9.0.post0 +pypdf==5.2.0 python-dotenv==1.0.1 +regex==2024.11.6 requests==2.32.3 scikit-learn==1.6.1 scipy==1.15.1 @@ -37,16 +38,10 @@ tabulate==0.9.0 threadpoolctl==3.5.0 tiktoken==0.8.0 tokenizers==0.21.0 -tqdm==4.67.1 -urllib3==2.4.0 -uvicorn==0.34.0 -websockets>=13.1.0 +uvicorn==0.22.0 +websockets==14.2 +Werkzeug==3.1.3 wheel==0.45.1 -xgboost-cpu==3.0.2 -bokeh==3.7.3 -pymc==5.23.0 -lightgbm==4.6.0 -arviz==0.21.0 -optuna==4.3.0 -litellm[proxy] -duckdb==1.3.2 \ No newline at end of file +gunicorn +psycopg2 +chardet \ No newline at end of file diff --git a/sample_output.txt b/sample_output.txt new file mode 100644 index 0000000000000000000000000000000000000000..76213181ad276466c32592b9ead42fed3dcb6a29 --- /dev/null +++ b/sample_output.txt @@ -0,0 +1,72 @@ + + +=== ERROR IN PLANNER_PREPROCESSING_AGENT === +Error in planner_preprocessing_agent: cannot do positional indexing on RangeIndex with these indexers [33.0] of type float +Traceback (most recent call last): + File "D:\Firebird2\Auto-Analyst-CS\auto-analystbackend\scripts\format_response.py", line 196, in execute_code_from_markdown + exec(block_code, context) # Execute the block + +Problem at this location: +1: +2: >>> print(df.head(33.)) <<< +3: +4: # planner_preprocessing_agent code end + +Full error details: + File "D:\Firebird2\Auto-Analyst-CS\auto-analyst-backend\venv\Lib\site-packages\pandas\core\indexes\base.py", line 4301, in _raise_invalid_indexer + raise TypeError(msg) +TypeError: cannot do positional indexing on RangeIndex with these indexers [33.0] of type float + + +=== ERROR IN PREPROCESSING_AGENT === +Error in preprocessing_agent: cannot do positional indexing on RangeIndex with these indexers [33.33] of type float +Traceback (most recent call last): + File "D:\Firebird2\Auto-Analyst-CS\auto-analyst-backend\scripts\format_response.py", line 196, in execute_code_from_markdown + exec(block_code, context) # Execute the block + +Problem at this location: +1: +2: >>> print(df.head(33.33)) <<< +3: +4: # preprocessing_agent code end + +Full error details: + File "D:\Firebird2\Auto-Analyst-CS\auto-analyst-backend\venv\Lib\site-packages\pandas\core\indexes\base.py", line 4301, in _raise_invalid_indexer + raise TypeError(msg) +TypeError: cannot do positional indexing on RangeIndex with these indexers [33.33] of type float + + +=== ERROR IN PLANNER_AGENT === +Error in planner_agent: NDFrame.head() takes from 1 to 2 positional arguments but 3 were given +Traceback (most recent call last): + File "D:\Firebird2\Auto-Analyst-CS\auto-analyst-backend\scripts\format_response.py", line 196, in execute_code_from_markdown + exec(block_code, context) # Execute the block + +Problem at this location: +1: +2: >>> print(df.head(32,3)) <<< +3: +4: # planner_agent code end + +Full error details: + ~~~~^^^^^^^^^^^^^^^^^^^^^ + File "", line 2, in +TypeError: NDFrame.head() takes from 1 to 2 positional arguments but 3 were given + + +=== ERROR IN STATS_AGENT === +Error in stats_agent: 'tuple' object is not callable +Traceback (most recent call last): + File "D:\Firebird2\Auto-Analyst-CS\auto-analyst-backend\scripts\format_response.py", line 196, in execute_code_from_markdown + exec(block_code, context) # Execute the block + +Problem at this location: +1: +2: >>> print(df.shape(3)) <<< +3: +4: # stats_agent code end + +Full error details: + ~~~~^^^^^^^^^^^^^^^^^^^^^ + File "", line 2, in +TypeError: 'tuple' object is not callable \ No newline at end of file diff --git a/scripts/create_test_user.py b/scripts/create_test_user.py new file mode 100644 index 0000000000000000000000000000000000000000..9951f80c2dcffec0d6b98c67c87cddf41f9935f5 --- /dev/null +++ b/scripts/create_test_user.py @@ -0,0 +1,45 @@ +import sys +import os + +# Add parent directory to path so we can import modules +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from init_db import User, session_factory + +def create_test_users(): + """Create test users for development""" + session = session_factory() + + try: + # Create a few test users + test_users = [ + {"username": "testuser1", "email": "user1@example.com"}, + {"username": "testuser2", "email": "user2@example.com"}, + {"username": "admin", "email": "admin@example.com"}, + {"username": "demo", "email": "demo@example.com"}, + {"username": "guest", "email": "guest@example.com"}, + ] + + for user_data in test_users: + # Check if user already exists + existing = session.query(User).filter(User.email == user_data["email"]).first() + if not existing: + user = User(**user_data) + session.add(user) + print(f"Created user: {user_data['username']}") + else: + print(f"User {user_data['username']} already exists") + + session.commit() + print("Test users created successfully") + + except Exception as e: + session.rollback() + print(f"Error creating test users: {e}") + + finally: + session.close() + +if __name__ == "__main__": + create_test_users() + print("Done!") \ No newline at end of file diff --git a/scripts/format_response.py b/scripts/format_response.py index 7902e45a4bd7f5df6368d2a4fe5ef73d5231a4c4..332bc6223a33f58eb0a560f2508917dee2308f82 100644 --- a/scripts/format_response.py +++ b/scripts/format_response.py @@ -8,7 +8,7 @@ import logging from src.utils.logger import Logger import textwrap -logger = Logger(__name__, level=logging.INFO, see_time=False, console_log=False) +logger = Logger(__name__, level="INFO", see_time=False, console_log=False) @contextlib.contextmanager def stdoutIO(stdout=None): @@ -44,33 +44,20 @@ API_KEY_PATTERNS = [ ] # Network request patterns -NETWORK_REQUEST_PATTERNS = re.compile(r"(requests\.|urllib\.|http\.client|httpx\.|socket\.connect\()") +NETWORK_REQUEST_PATTERNS = re.compile(r"(requests\.|urllib\.|http\.|\.post\(|\.get\(|\.connect\()") -# DataFrame creation with hardcoded data - block only this specific pattern - -def check_security_concerns(code_str, dataset_names): +def check_security_concerns(code_str): """Check code for security concerns and return info about what was found""" security_concerns = { "has_concern": False, + "messages": [], "blocked_imports": False, "blocked_dynamic_imports": False, "blocked_env_access": False, "blocked_file_access": False, "blocked_api_keys": False, - "blocked_network": False, - "blocked_dataframe_invention": False, # Add this new field - "messages": [] + "blocked_network": False } - - dataset_names_pattern = "|".join(re.escape(name) for name in dataset_names) - DATAFRAME_INVENTION_PATTERN = re.compile( - rf"({dataset_names_pattern})\s*=\s*pd\.DataFrame\s*\(\s*\{{\s*[^}}]*\}}", - re.MULTILINE - ) - if DATAFRAME_INVENTION_PATTERN.search(code_str): - security_concerns["has_concern"] = True - security_concerns["blocked_dataframe_invention"] = True - security_concerns["messages"].append(f"DataFrame creation blocked for dataset variables: {', '.join(dataset_names)}") # Check for sensitive imports if IMPORT_PATTERN.search(code_str) or FROM_IMPORT_PATTERN.search(code_str): @@ -110,19 +97,11 @@ def check_security_concerns(code_str, dataset_names): security_concerns["blocked_network"] = True security_concerns["messages"].append("Network requests blocked") - - return security_concerns -def clean_code_for_security(code_str, security_concerns, dataset_names): +def clean_code_for_security(code_str, security_concerns): """Apply security modifications to the code based on detected concerns""" - modified_code = code_str - dataset_names_pattern = "|".join(re.escape(name) for name in dataset_names) - DATAFRAME_INVENTION_PATTERN = re.compile( - rf"({dataset_names_pattern})\s*=\s*pd\.DataFrame\s*\(\s*\{{\s*[^}}]*\}}", - re.MULTILINE - ) # Block sensitive imports if needed if security_concerns["blocked_imports"]: @@ -150,13 +129,6 @@ def clean_code_for_security(code_str, security_concerns, dataset_names): if security_concerns["blocked_network"]: modified_code = NETWORK_REQUEST_PATTERNS.sub(r'"BLOCKED_NETWORK_REQUEST"', modified_code) - # Block DataFrame creation with hardcoded data if detected - if security_concerns["blocked_dataframe_invention"]: - modified_code = DATAFRAME_INVENTION_PATTERN.sub( - r"# BLOCKED_DATAFRAME_INVENTION: \g<0>", - modified_code - ) - # Add warning banner if needed if security_concerns["has_concern"]: security_message = "⚠️ SECURITY WARNING: " + ". ".join(security_concerns["messages"]) + "." @@ -164,104 +136,6 @@ def clean_code_for_security(code_str, security_concerns, dataset_names): return modified_code -def format_correlation_output(text): - """Format correlation matrix output for better readability""" - lines = text.split('\n') - formatted_lines = [] - - for line in lines: - # Skip empty lines at the beginning - if not line.strip() and not formatted_lines: - continue - - if not line.strip(): - formatted_lines.append(line) - continue - - # Check if this line contains correlation values or variable names - stripped_line = line.strip() - parts = stripped_line.split() - - if len(parts) > 1: - # Check if this is a header line with variable names - if all(part.replace('_', '').replace('-', '').isalpha() for part in parts): - # This is a header row with variable names - formatted_header = f"{'':12}" # Empty first column for row labels - for part in parts: - formatted_header += f"{part:>12}" - formatted_lines.append(formatted_header) - elif any(char.isdigit() for char in stripped_line) and ('.' in stripped_line or '-' in stripped_line): - # This looks like a correlation line with numbers - row_name = parts[0] if parts else "" - values = parts[1:] if len(parts) > 1 else [] - - formatted_row = f"{row_name:<12}" - for value in values: - try: - val = float(value) - formatted_row += f"{val:>12.3f}" - except ValueError: - formatted_row += f"{value:>12}" - - formatted_lines.append(formatted_row) - else: - # Other lines (like titles) - formatted_lines.append(line) - else: - formatted_lines.append(line) - - return '\n'.join(formatted_lines) - -def format_summary_stats(text): - """Format summary statistics for better readability""" - lines = text.split('\n') - formatted_lines = [] - - for line in lines: - if not line.strip(): - formatted_lines.append(line) - continue - - # Check if this is a header line with statistical terms only (missing first column) - stripped_line = line.strip() - if any(stat in stripped_line.lower() for stat in ['count', 'mean', 'median', 'std', 'min', 'max', '25%', '50%', '75%']): - parts = stripped_line.split() - # Check if this is a header row (starts with statistical terms) - if parts and parts[0].lower() in ['count', 'mean', 'median', 'std', 'min', 'max', '25%', '50%', '75%']: - # This is a header row - add proper spacing - formatted_header = f"{'':12}" # Empty first column for row labels - for part in parts: - formatted_header += f"{part:>15}" - formatted_lines.append(formatted_header) - else: - # This is a data row - format normally - row_name = parts[0] if parts else "" - values = parts[1:] if len(parts) > 1 else [] - - formatted_row = f"{row_name:<12}" - for value in values: - try: - if '.' in value or 'e' in value.lower(): - val = float(value) - if abs(val) >= 1000000: - formatted_row += f"{val:>15.2e}" - elif abs(val) >= 1: - formatted_row += f"{val:>15.2f}" - else: - formatted_row += f"{val:>15.6f}" - else: - val = int(value) - formatted_row += f"{val:>15}" - except ValueError: - formatted_row += f"{value:>15}" - - formatted_lines.append(formatted_row) - else: - # Other lines (titles, etc.) - keep as is - formatted_lines.append(line) - - return '\n'.join(formatted_lines) - def clean_print_statements(code_block): """ This function cleans up any `print()` statements that might contain unwanted `\n` characters. @@ -304,14 +178,6 @@ def format_code_block(code_str): return f'\n{code_clean}\n' def format_code_backticked_block(code_str): - # Add None check at the beginning - if code_str is None: - return "```python\n# No code available\n```" # Return empty code block instead of None - - # Add type check to ensure it's a string - if not isinstance(code_str, str): - return f"```python\n# Invalid code type: {type(code_str)}\n```" - code_clean = re.sub(r'^```python\n?', '', code_str, flags=re.MULTILINE) code_clean = re.sub(r'\n```$', '', code_clean) # Only match assignments at top level (not indented) @@ -320,7 +186,6 @@ def format_code_backticked_block(code_str): # Remove reading the csv file if it's already in the context modified_code = re.sub(r"df\s*=\s*pd\.read_csv\([\"\'].*?[\"\']\).*?(\n|$)", '', code_clean) - modified_code = re.sub(r'^(\s*)(df\s*=.*)$', r'\1# \2', code_clean, flags=re.MULTILINE) # Only match assignments at top level (not indented) # 1. Remove 'df = pd.DataFrame()' if it's at the top level @@ -344,7 +209,8 @@ def format_code_backticked_block(code_str): return f'```python\n{code_clean}\n```' -def execute_code_from_markdown(code_str, datasets=None): +# In format_response.py, modify the execute_code function: +def execute_code_from_markdown(code_str, dataframe=None): import pandas as pd import plotly.express as px import plotly @@ -355,169 +221,13 @@ def execute_code_from_markdown(code_str, datasets=None): import re import traceback import sys - from io import StringIO, BytesIO - import base64 + from io import StringIO - context_names = list(datasets.keys()) # Check for security concerns in the code - security_concerns = check_security_concerns(code_str, context_names) + security_concerns = check_security_concerns(code_str) # Apply security modifications to the code - modified_code = clean_code_for_security(code_str, security_concerns, context_names) - - # Enhanced print function that detects and formats tabular data - captured_outputs = [] - original_print = print - - # Set pandas display options for full table display - pd.set_option('display.max_columns', None) - pd.set_option('display.max_rows', 20) # Limit to 20 rows instead of unlimited - pd.set_option('display.width', None) - pd.set_option('display.max_colwidth', 50) - pd.set_option('display.expand_frame_repr', False) - - - - def enhanced_print(*args, **kwargs): - # Convert all args to strings - str_args = [str(arg) for arg in args] - output_text = kwargs.get('sep', ' ').join(str_args) - - # Special case for DataFrames - use pipe delimiter and clean format - if isinstance(args[0], pd.DataFrame) and len(args) == 1: - # Format DataFrame with pipe delimiter using to_csv for reliable column separation - df = args[0] - - # Use StringIO to capture CSV output with pipe delimiter - from io import StringIO - csv_buffer = StringIO() - - # Export to CSV with pipe delimiter, preserving index - df.to_csv(csv_buffer, sep='|', index=True, float_format='%.6g') - csv_output = csv_buffer.getvalue() - - # Clean up the CSV output - remove quotes and extra formatting - lines = csv_output.strip().split('\n') - cleaned_lines = [] - - for line in lines: - # Remove any quotes that might have been added by to_csv - clean_line = line.replace('"', '') - # Split by pipe, strip whitespace from each part, then rejoin - parts = [part.strip() for part in clean_line.split('|')] - cleaned_lines.append(' | '.join(parts)) - - output_text = '\n'.join(cleaned_lines) - captured_outputs.append(f"\n{output_text}\n") - original_print(output_text) - return - - # Detect if this looks like tabular data (generic approach) - is_table = False - - # Check for table patterns: - # 1. Multiple lines with consistent spacing - lines = output_text.split('\n') - if len(lines) > 2: - # Count lines that look like they have multiple columns (2+ spaces between words) - multi_column_lines = sum(1 for line in lines if len(line.split()) > 1 and ' ' in line) - if multi_column_lines >= 2: # At least 2 lines with multiple columns - is_table = True - - # Check for pandas DataFrame patterns like index with column names - if any(re.search(r'^\s*\d+\s+', line) for line in lines): - # Look for lines starting with an index number followed by spaces - is_table = True - - # Look for table-like structured output with multiple rows of similar format - if len(lines) >= 3: - # Sample a few lines to check for consistent structure - sample_lines = [lines[i] for i in range(min(len(lines), 5)) if i < len(lines) and lines[i].strip()] - - # Check for consistent whitespace patterns - if len(sample_lines) >= 2: - # Get positions of whitespace groups in first line - whitespace_positions = [] - for i, line in enumerate(sample_lines): - if not line.strip(): - continue - positions = [m.start() for m in re.finditer(r'\s{2,}', line)] - if i == 0: - whitespace_positions = positions - elif len(positions) == len(whitespace_positions): - # Check if whitespace positions are roughly the same - is_similar = all(abs(pos - whitespace_positions[j]) <= 3 - for j, pos in enumerate(positions) - if j < len(whitespace_positions)) - if is_similar: - is_table = True - - # 2. Contains common table indicators - if any(indicator in output_text.lower() for indicator in [ - 'count', 'mean', 'std', 'min', 'max', '25%', '50%', '75%', # Summary stats - 'correlation', 'corr', # Correlation tables - 'coefficient', 'r-squared', 'p-value', # Regression tables - ]): - is_table = True - - # 3. Has many decimal numbers (likely a data table) - if output_text.count('.') > 5 and len(lines) > 2: - is_table = True - - # If we have detected a table, convert space-delimited to pipe-delimited format - if is_table: - # Convert the table to pipe-delimited format for better parsing in frontend - formatted_lines = [] - for line in lines: - if not line.strip(): - formatted_lines.append(line) # Keep empty lines - continue - - # Split by multiple spaces and join with pipe delimiter - parts = re.split(r'\s{2,}', line.strip()) - if parts: - formatted_lines.append(" | ".join(parts)) - else: - formatted_lines.append(line) - - # Use the pipe-delimited format - output_text = "\n".join(formatted_lines) - - # Format and mark the output for table processing in UI - captured_outputs.append(f"\n{output_text}\n") - else: - captured_outputs.append(output_text) - - # Also use original print for stdout capture - original_print(*args, **kwargs) - - # Custom matplotlib capture function - def capture_matplotlib_chart(): - """Capture current matplotlib figure as base64 encoded image""" - try: - fig = plt.gcf() # Get current figure - if fig.get_axes(): # Check if figure has any plots - buffer = BytesIO() - fig.savefig(buffer, format='png', dpi=150, bbox_inches='tight', - facecolor='white', edgecolor='none') - buffer.seek(0) - img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8') - buffer.close() - plt.close(fig) # Close the figure to free memory - return img_base64 - return None - except Exception: - return None - - # Store original plt.show function - original_plt_show = plt.show - - def custom_plt_show(*args, **kwargs): - """Custom plt.show that captures the chart instead of displaying it""" - img_base64 = capture_matplotlib_chart() - if img_base64: - matplotlib_outputs.append(img_base64) - # Don't call original show to prevent display + modified_code = clean_code_for_security(code_str, security_concerns) context = { 'pd': pd, @@ -529,18 +239,8 @@ def execute_code_from_markdown(code_str, datasets=None): '__import__': __import__, 'sns': sns, 'np': np, - 'json_outputs': [], # List to store multiple Plotly JSON outputs - 'matplotlib_outputs': [], # List to store matplotlib chart images as base64 - 'print': enhanced_print # Replace print with our enhanced version + 'json_outputs': [] # List to store multiple Plotly JSON outputs } - - # Add matplotlib_outputs to local scope for the custom show function - matplotlib_outputs = context['matplotlib_outputs'] - - # Replace plt.show with our custom function - plt.show = custom_plt_show - - # Modify code to store multiple JSON outputs modified_code = re.sub( @@ -553,7 +253,8 @@ def execute_code_from_markdown(code_str, datasets=None): r'(\w*_?)fig(\w*)\.to_html\(.*?\)', r'json_outputs.append(plotly.io.to_json(\1fig\2, pretty=True))', modified_code - ) + ) + # Remove reading the csv file if it's already in the context modified_code = re.sub(r"df\s*=\s*pd\.read_csv\([\"\'].*?[\"\']\).*?(\n|$)", '', modified_code) @@ -565,70 +266,27 @@ def execute_code_from_markdown(code_str, datasets=None): modified_code, flags=re.MULTILINE ) - - - # Custom display function for DataFrames to show head + tail for large datasets - original_repr = pd.DataFrame.__repr__ - - def custom_df_repr(self): - if len(self) > 15: - # For large DataFrames, show first 10 and last 5 rows - head_part = self.head(10) - tail_part = self.tail(5) - - head_str = head_part.__repr__() - tail_str = tail_part.__repr__() - - # Extract just the data rows (skip the header from tail) - tail_lines = tail_str.split('\n') - tail_data = '\n'.join(tail_lines[1:]) # Skip header line - - return f"{head_str}\n...\n{tail_data}" - else: - return original_repr(self) - - # Apply custom representation temporarily - pd.DataFrame.__repr__ = custom_df_repr # If a dataframe is provided, add it to the context - for dataset_name, dataset_df in datasets.items(): - if dataset_df is not None: - context[dataset_name] = dataset_df - logger.log_message(f"Added dataset '{dataset_name}' to execution context", level=logging.DEBUG) - + if dataframe is not None: + context['df'] = dataframe # remove pd.read_csv() if it's already in the context modified_code = re.sub(r"pd\.read_csv\(\s*[\"\'].*?[\"\']\s*\)", '', modified_code) # Remove sample dataframe lines with multiple array values modified_code = re.sub(r"^# Sample DataFrames?.*?(\n|$)", '', modified_code, flags=re.MULTILINE | re.IGNORECASE) - - # Replace plt.savefig() calls with plt.show() to ensure plots are displayed - modified_code = re.sub(r'plt\.savefig\([^)]*\)', 'plt.show()', modified_code) - - # Instead of removing plt.show(), keep them - they'll be handled by our custom function - # Also handle seaborn plots that might not have explicit plt.show() - # Add plt.show() after seaborn plot functions if not already present - seaborn_plot_functions = [ - 'sns.scatterplot', 'sns.lineplot', 'sns.barplot', 'sns.boxplot', 'sns.violinplot', - 'sns.stripplot', 'sns.swarmplot', 'sns.pointplot', 'sns.catplot', 'sns.relplot', - 'sns.displot', 'sns.histplot', 'sns.kdeplot', 'sns.ecdfplot', 'sns.rugplot', - 'sns.distplot', 'sns.jointplot', 'sns.pairplot', 'sns.FacetGrid', 'sns.PairGrid', - 'sns.heatmap', 'sns.clustermap', 'sns.regplot', 'sns.lmplot', 'sns.residplot' - ] - - # Add automatic plt.show() after seaborn plots if not already present - for func in seaborn_plot_functions: - pattern = rf'({re.escape(func)}\([^)]*\)(?:\.[^(]*\([^)]*\))*)' - def add_show(match): - plot_call = match.group(1) - # Check if the next non-empty line already has plt.show() - return f'{plot_call}\nplt.show()' - modified_code = re.sub(pattern, add_show, modified_code) + # Remove plt.show() statements + modified_code = re.sub(r"plt\.show\(\).*?(\n|$)", '', modified_code) # Only add df = pd.read_csv() if no dataframe was provided and the code contains pd.read_csv - + if dataframe is None and 'pd.read_csv' not in modified_code: + modified_code = re.sub( + r'import pandas as pd', + r'import pandas as pd\n\n# Read Housing.csv\ndf = pd.read_csv("Housing.csv")', + modified_code + ) # Identify code blocks by comments code_blocks = [] @@ -656,42 +314,11 @@ def execute_code_from_markdown(code_str, datasets=None): all_outputs = [] for block_name, block_code in code_blocks: try: - # Clear captured outputs for each block - captured_outputs.clear() - - # Fix indentation issues before execution - try: - block_code = textwrap.dedent(block_code) - except Exception as dedent_error: - logger.log_message(f"Failed to dedent code block '{block_name}': {str(dedent_error)}", level=logging.WARNING) - with stdoutIO() as s: exec(block_code, context) # Execute the block - - # Get both stdout and our enhanced captured outputs - stdout_output = s.getvalue() - - # Combine outputs, preferring our enhanced format when available - if captured_outputs: - combined_output = '\n'.join(captured_outputs) - else: - combined_output = stdout_output - - all_outputs.append((block_name, combined_output, None)) # None means no error + output = s.getvalue() + all_outputs.append((block_name, output, None)) # None means no error except Exception as e: - # Reset pandas options in case of error - pd.reset_option('display.max_columns') - pd.reset_option('display.max_rows') - pd.reset_option('display.width') - pd.reset_option('display.max_colwidth') - pd.reset_option('display.expand_frame_repr') - - # Restore original DataFrame representation in case of error - pd.DataFrame.__repr__ = original_repr - - # Restore original plt.show - plt.show = original_plt_show - error_traceback = traceback.format_exc() # Extract error message and error type @@ -838,23 +465,9 @@ def execute_code_from_markdown(code_str, datasets=None): all_outputs.append((block_name, None, formatted_error)) - # Reset pandas options after execution - pd.reset_option('display.max_columns') - pd.reset_option('display.max_rows') - pd.reset_option('display.width') - pd.reset_option('display.max_colwidth') - pd.reset_option('display.expand_frame_repr') - - # Restore original DataFrame representation - pd.DataFrame.__repr__ = original_repr - - # Restore original plt.show - plt.show = original_plt_show - # Compile all outputs and errors output_text = "" json_outputs = context.get('json_outputs', []) - matplotlib_outputs = context.get('matplotlib_outputs', []) error_found = False for block_name, output, error in all_outputs: @@ -865,9 +478,9 @@ def execute_code_from_markdown(code_str, datasets=None): output_text += f"\n\n=== OUTPUT FROM {block_name.upper()}_AGENT ===\n{output}\n" if error_found: - return output_text, [], [] + return output_text, [] else: - return output_text, json_outputs, matplotlib_outputs + return output_text, json_outputs def format_plan_instructions(plan_instructions): @@ -875,127 +488,54 @@ def format_plan_instructions(plan_instructions): Format any plan instructions (JSON string or dict) into markdown sections per agent. """ # Parse input into a dict - - if "basic_qa_agent" in str(plan_instructions): - return "**Non-Data Request**: Please ask a data related query, don't waste credits!" - - try: if isinstance(plan_instructions, str): - try: - instructions = json.loads(plan_instructions) - except json.JSONDecodeError as e: - # Try to clean the string if it's not valid JSON - cleaned_str = plan_instructions.strip() - if cleaned_str.startswith("'") and cleaned_str.endswith("'"): - cleaned_str = cleaned_str[1:-1] - try: - instructions = json.loads(cleaned_str) - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON format in plan instructions: {str(e)}") + instructions = json.loads(plan_instructions) elif isinstance(plan_instructions, dict): instructions = plan_instructions else: - raise TypeError(f"Unsupported plan instructions type: {type(plan_instructions)}") - except Exception as e: - raise ValueError(f"Error processing plan instructions: {str(e)} + {dspy.settings.lm} ") + return f"Unsupported plan instructions type: {type(plan_instructions)}" + except json.JSONDecodeError: + return f"Invalid plan instructions: {plan_instructions}" # logger.log_message(f"Plan instructions: {instructions}", level=logging.INFO) - - markdown_lines = [] for agent, content in instructions.items(): - if agent != 'basic_qa_agent': - agent_title = agent.replace('_', ' ').title() - markdown_lines.append(f"#### {agent_title}") - if isinstance(content, dict): - # Handle 'create' key - create_vals = content.get('create', []) - if create_vals: - markdown_lines.append(f"- **Create**:") - for item in create_vals: - markdown_lines.append(f" - {item}") - else: - markdown_lines.append(f"- **Create**: None") - - # Handle 'use' key - use_vals = content.get('use', []) - if use_vals: - markdown_lines.append(f"- **Use**:") - for item in use_vals: - markdown_lines.append(f" - {item}") - else: - markdown_lines.append(f"- **Use**: None") - - # Handle 'instruction' key - instr = content.get('instruction') - if isinstance(instr, str) and instr: - markdown_lines.append(f"- **Instruction**: {instr}") - else: - markdown_lines.append(f"- **Instruction**: None") + agent_title = agent.replace('_', ' ').title() + markdown_lines.append(f"#### {agent_title}") + if isinstance(content, dict): + # Handle 'create' key + create_vals = content.get('create', []) + if create_vals: + markdown_lines.append(f"- **Create**:") + for item in create_vals: + markdown_lines.append(f" - {item}") else: - # Fallback for non-dict content - markdown_lines.append(f"- {content}") - markdown_lines.append("") # blank line between agents - else: - markdown_lines.append(f"**Non-Data Request**: {content.get('instruction')}") - - return "\n".join(markdown_lines).strip() - - # Return empty string if no complexity found - return "" + markdown_lines.append(f"- **Create**: None") + + # Handle 'use' key + use_vals = content.get('use', []) + if use_vals: + markdown_lines.append(f"- **Use**:") + for item in use_vals: + markdown_lines.append(f" - {item}") + else: + markdown_lines.append(f"- **Use**: None") -def format_complexity(instructions): - markdown_lines = [] - complexity = None # Initialize complexity to avoid UnboundLocalError - - # Extract complexity from various possible locations in the structure - if isinstance(instructions, dict): - # Case 1: Direct complexity field - if 'complexity' in instructions: - complexity = instructions['complexity'] - # Case 2: Complexity in 'plan' object - elif 'plan' in instructions and isinstance(instructions['plan'], dict): - if 'complexity' in instructions['plan']: - complexity = instructions['plan']['complexity'] + # Handle 'instruction' key + instr = content.get('instruction') + if isinstance(instr, str) and instr: + markdown_lines.append(f"- **Instruction**: {instr}") + else: + markdown_lines.append(f"- **Instruction**: None") else: - complexity = "unrelated" - - # Override for basic_qa_agent cases - if 'plan' in instructions and isinstance(instructions['plan'], str) and "basic_qa_agent" in instructions['plan']: - complexity = "unrelated" - else: - # If instructions is not a dict, default to unrelated - complexity = "unrelated" - - if complexity: - # Pink color scheme variations - color_map = { - "unrelated": "#FFB6B6", # Light pink - "basic": "#FF9E9E", # Medium pink - "intermediate": "#FF7F7F", # Main pink - "advanced": "#FF5F5F" # Dark pink - } - - indicator_map = { - "unrelated": "○", - "basic": "●", - "intermediate": "●●", - "advanced": "●●●" - } - - color = color_map.get(complexity.lower(), "#FFB6B6") # Default to light pink - indicator = indicator_map.get(complexity.lower(), "○") - - # Slightly larger display with pink styling - markdown_lines.append(f"
{indicator} {complexity}
\n") + # Fallback for non-dict content + markdown_lines.append(f"- {content}") + markdown_lines.append("") # blank line between agents - return "\n".join(markdown_lines).strip() + return "\n".join(markdown_lines).strip() - # Return empty string if no complexity found - return "" - -def format_response_to_markdown(api_response, agent_name = None, datasets=None): +def format_response_to_markdown(api_response, agent_name = None, dataframe=None): try: markdown = [] # logger.log_message(f"API response for {agent_name} at {time.strftime('%Y-%m-%d %H:%M:%S')}: {api_response}", level=logging.INFO) @@ -1020,81 +560,64 @@ def format_response_to_markdown(api_response, agent_name = None, datasets=None): if "memory" in agent or not content: continue - if "complexity" in content: - complexity_result = format_complexity(content) - if complexity_result: # Only append if not empty - markdown.append(f"{complexity_result}") - markdown.append(f"\n## {agent.replace('_', ' ').title()}\n") if agent == "analytical_planner": - logger.log_message(f"Analytical planner content: {content}", level=logging.INFO) - if 'plan_desc' in content and content['plan_desc']: - markdown.append(f"### Reasoning {content['plan_desc']}") + # logger.log_message(f"Analytical planner content: {content}", level=logging.INFO) + if 'plan_desc' in content: + markdown.append(f"### Reasoning\n{content['plan_desc']}\n") if 'plan_instructions' in content: - plan_result = format_plan_instructions(content['plan_instructions']) - if plan_result: # Only append if not empty/None - markdown.append(f"{plan_result}") + markdown.append(f"### Plan Instructions\n{format_plan_instructions(content['plan_instructions'])}\n") else: - if content.get('rationale'): - markdown.append(f"### Reasoning {content['rationale']}") + markdown.append(f"### Reasoning\n{content['rationale']}\n") else: if "rationale" in content: - if content.get('rationale'): - markdown.append(f"### Reasoning{content['rationale']}") + markdown.append(f"### Reasoning\n{content['rationale']}\n") + + if 'code' in content: + markdown.append(f"### Code Implementation\n{format_code_backticked_block(content['code'])}\n") + # if agent_name is not None: + # # execute the code + # clean_code = format_code_block(content['code']) + # output, json_outputs = execute_code_from_markdown(clean_code, dataframe) + # if output: + # markdown.append("### Execution Output\n") + # markdown.append(f"```output\n{output}\n```\n") + + # if json_outputs: + # markdown.append("### Plotly JSON Outputs\n") + # for idx, json_output in enumerate(json_outputs): + # if len(json_output) > 1000000: # If JSON is larger than 1MB + # logger.log_message(f"Large JSON output detected: {len(json_output)} bytes", level=logging.WARNING) + # markdown.append(f"```plotly\n{json_output}\n```\n") - if 'code' in content and content['code'] is not None: - formatted_code = format_code_backticked_block(content['code']) - if formatted_code: # Check if formatted_code is not None - markdown.append(f"### Code Implementation\n{formatted_code}\n") - if 'answer' in content and content['answer']: - markdown.append(f"### Answer{content['answer']} Please ask a query about the data") if 'summary' in content: - import re - summary_text = content['summary'] - summary_text = re.sub(r'```python\n(.*?)\n```', '', summary_text, flags=re.DOTALL) - + # make the summary a bullet-point list + summary_lines = remove_code_block_from_summary(content['summary']) + summary_lines = content['summary'].split('\n') + # remove code block from summary markdown.append("### Summary\n") - - # Extract pre-list intro, bullet points, and post-list text - intro_match = re.split(r'\(\d+\)', summary_text, maxsplit=1) - if len(intro_match) > 1: - intro_text = intro_match[0].strip() - rest_text = "(1)" + intro_match[1] # reattach for bullet parsing - else: - intro_text = summary_text.strip() - rest_text = "" - - if intro_text: - markdown.append(f"{intro_text}\n") - - # Split bullets at numbered items like (1)...(8) - bullets = re.split(r'\(\d+\)', rest_text) - bullets = [b.strip(" ,.\n") for b in bullets if b.strip()] - - # Check for post-list content (anything after the last number) - for i, bullet in enumerate(bullets): - markdown.append(f"* {bullet}\n") - - - + for line in summary_lines: + if line != "": + if line.strip().startswith('•') or line.strip().startswith('-') or line.strip().startswith('*'): + line = line.strip().replace('•', '').replace('-', '').replace('*', '') + markdown.append(f"* {line.strip()}\n") + else: + markdown.append(f"{line.strip()}\n") if 'refined_complete_code' in content and 'summary' in content: try: if content['refined_complete_code'] is not None and content['refined_complete_code'] != "": clean_code = format_code_block(content['refined_complete_code']) markdown_code = format_code_backticked_block(content['refined_complete_code']) - output, json_outputs, matplotlib_outputs = execute_code_from_markdown(clean_code, datasets) + output, json_outputs = execute_code_from_markdown(clean_code, dataframe) elif "```python" in content['summary']: clean_code = format_code_block(content['summary']) markdown_code = format_code_backticked_block(content['summary']) - output, json_outputs, matplotlib_outputs = execute_code_from_markdown(clean_code, datasets) + output, json_outputs = execute_code_from_markdown(clean_code, dataframe) except Exception as e: logger.log_message(f"Error in execute_code_from_markdown: {str(e)}", level=logging.ERROR) markdown_code = f"**Error**: {str(e)}" - output = None - json_outputs = [] - matplotlib_outputs = [] # continue if markdown_code is not None: @@ -1108,30 +631,45 @@ def format_response_to_markdown(api_response, agent_name = None, datasets=None): markdown.append("### Plotly JSON Outputs\n") for idx, json_output in enumerate(json_outputs): markdown.append(f"```plotly\n{json_output}\n```\n") - - if matplotlib_outputs: - markdown.append("### Matplotlib/Seaborn Charts\n") - for idx, img_base64 in enumerate(matplotlib_outputs): - markdown.append(f"```matplotlib\n{img_base64}\n```\n") # if agent_name is not None: # if f"memory_{agent_name}" in api_response: # markdown.append(f"### Memory\n{api_response[f'memory_{agent_name}']}\n") except Exception as e: logger.log_message(f"Error in format_response_to_markdown: {str(e)}", level=logging.ERROR) - return f"error formating markdown {str(e)}" + return f"{str(e)}" # logger.log_message(f"Generated markdown content for agent '{agent_name}' at {time.strftime('%Y-%m-%d %H:%M:%S')}: {markdown}, length: {len(markdown)}", level=logging.INFO) if not markdown or len(markdown) <= 1: - logger.log_message( - f"Invalid markdown content for agent '{agent_name}' at {time.strftime('%Y-%m-%d %H:%M:%S')}: " - f"Content: '{markdown}', Type: {type(markdown)}, Length: {len(markdown) if markdown else 0}, " - f"API Response: {api_response}", - level=logging.ERROR - ) - return "" + logger.log_message(f"Generated markdown (ERROR) content for agent '{agent_name}' at {time.strftime('%Y-%m-%d %H:%M:%S')}: {markdown}, length: {len(markdown)}, api_response: {api_response}", level=logging.INFO) + return "Please provide a valid query..." return '\n'.join(markdown) +# Example usage with dummy data +if __name__ == "__main__": + sample_response = { + "code_combiner_agent": { + "reasoning": "Sample reasoning for multiple charts.", + "refined_complete_code": """ +```python +import plotly.express as px +import pandas as pd + +# Sample Data +df = pd.DataFrame({'Category': ['A', 'B', 'C'], 'Values': [10, 20, 30]}) + +# First Chart +fig = px.bar(df, x='Category', y='Values', title='Bar Chart') +fig.show() + +# Second Chart +fig2 = px.pie(df, values='Values', names='Category', title='Pie Chart') +fig2.show() +``` +""" + } + } + formatted_md = format_response_to_markdown(sample_response) \ No newline at end of file diff --git a/scripts/generate_test_data.py b/scripts/generate_test_data.py new file mode 100644 index 0000000000000000000000000000000000000000..4e9c62c54b86dc8ad9e42fcf23d3d047314eac0c --- /dev/null +++ b/scripts/generate_test_data.py @@ -0,0 +1,90 @@ +import sys +import os +import random +from datetime import datetime, timedelta +import sqlite3 + +# Add parent directory to path so we can import modules +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from init_db import ModelUsage, session_factory + +# Models and providers to use in test data +MODELS = { + "gpt-3.5-turbo": {"provider": "OpenAI", "cost_per_1k": 0.0015}, + "gpt-4": {"provider": "OpenAI", "cost_per_1k": 0.03}, + "gpt-4o": {"provider": "OpenAI", "cost_per_1k": 0.01}, + "gpt-4o-mini": {"provider": "OpenAI", "cost_per_1k": 0.0015}, + "o1-mini": {"provider": "OpenAI", "cost_per_1k": 0.00015}, + "claude-3-opus": {"provider": "Anthropic", "cost_per_1k": 0.015}, + "claude-3-sonnet": {"provider": "Anthropic", "cost_per_1k": 0.008}, + "claude-3-haiku": {"provider": "Anthropic", "cost_per_1k": 0.003}, + "llama-3-8b": {"provider": "Groq", "cost_per_1k": 0.0005}, + "llama-3-70b": {"provider": "Groq", "cost_per_1k": 0.002}, +} + +# User IDs to use (can be random if you don't have specific users) +USER_IDS = [1, 2, 3, 4, 5] + +def generate_test_data(num_records=100): + """Generate test model usage data""" + session = session_factory() + + try: + # Generate records for the past 30 days + end_date = datetime.utcnow() + start_date = end_date - timedelta(days=30) + + for _ in range(num_records): + # Random timestamp within the date range + random_days = random.randint(0, 30) + timestamp = end_date - timedelta(days=random_days, + hours=random.randint(0, 23), + minutes=random.randint(0, 59)) + + # Select random model and user + model_name = random.choice(list(MODELS.keys())) + model_info = MODELS[model_name] + user_id = random.choice(USER_IDS) + + # Generate random token counts + prompt_tokens = random.randint(100, 1000) + completion_tokens = random.randint(50, 500) + total_tokens = prompt_tokens + completion_tokens + + # Calculate cost + cost = (total_tokens / 1000) * model_info["cost_per_1k"] + + # Create model usage record + usage = ModelUsage( + user_id=user_id, + chat_id=random.randint(1, 50), # Random chat ID + model_name=model_name, + provider=model_info["provider"], + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + query_size=prompt_tokens * 4, # Approximate characters + response_size=completion_tokens * 4, # Approximate characters + cost=cost, + timestamp=timestamp, + is_streaming=random.choice([True, False]), + request_time_ms=random.randint(500, 5000) # Between 0.5 and 5 seconds + ) + session.add(usage) + + session.commit() + print(f"Successfully generated {num_records} test records") + + except Exception as e: + session.rollback() + print(f"Error generating test data: {e}") + + finally: + session.close() + +if __name__ == "__main__": + # Default to 100 records, but allow command line override + num_records = int(sys.argv[1]) if len(sys.argv) > 1 else 100 + generate_test_data(num_records) + print("Done! The model_usage table has been populated with test data.") \ No newline at end of file diff --git a/scripts/init_production_db.py b/scripts/init_production_db.py deleted file mode 100644 index 04cf091242580c6e35c820df3e6e1f159c0c134c..0000000000000000000000000000000000000000 --- a/scripts/init_production_db.py +++ /dev/null @@ -1,191 +0,0 @@ -#!/usr/bin/env python3 -""" -Production database initialization script. -This ensures templates are populated properly and verifies database health. -SAFE for PostgreSQL/RDS - only creates tables on SQLite databases. -""" - -import sys -import os -import logging -from datetime import datetime, UTC - -# Add the project root to the Python path -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -from src.db.init_db import init_db, session_factory, engine, is_postgres_db -from src.db.schemas.models import Base, AgentTemplate, UserTemplatePreference -from scripts.populate_agent_templates import populate_templates -from sqlalchemy import inspect, text -from src.utils.logger import Logger - -logger = Logger("init_production_db", see_time=True, console_log=True) - -def get_database_type(): - """Get the database type (sqlite or postgresql).""" - try: - if is_postgres_db(): - return "postgresql" - else: - return "sqlite" - except Exception as e: - logger.log_message(f"Error determining database type: {e}", logging.ERROR) - return "unknown" - -def check_table_exists(table_name: str) -> bool: - """Check if a table exists in the database.""" - try: - inspector = inspect(engine) - tables = inspector.get_table_names() - return table_name in tables - except Exception as e: - logger.log_message(f"Error checking table existence: {e}", logging.ERROR) - return False - -def verify_database_schema(): - """Verify that all required tables exist. Only create tables on SQLite.""" - db_type = get_database_type() - logger.log_message(f"🔍 Verifying database schema for {db_type.upper()} database...", logging.INFO) - - required_tables = [ - 'users', 'chats', 'messages', 'model_usage', 'code_executions', - 'message_feedback', 'deep_analysis_reports', 'agent_templates', - 'user_template_preferences' - ] - - missing_tables = [] - existing_tables = [] - - for table in required_tables: - if not check_table_exists(table): - missing_tables.append(table) - logger.log_message(f"❌ Missing table: {table}", logging.WARNING) - else: - existing_tables.append(table) - logger.log_message(f"✅ Table exists: {table}", logging.INFO) - - if missing_tables: - if db_type == "sqlite": - logger.log_message(f"🔧 Creating missing tables on SQLite: {missing_tables}", logging.INFO) - try: - # Safe to create tables on SQLite - Base.metadata.create_all(engine) - logger.log_message("✅ All tables created successfully on SQLite", logging.INFO) - except Exception as e: - logger.log_message(f"❌ Failed to create tables: {e}", logging.ERROR) - raise - else: - # PostgreSQL/RDS - DO NOT create tables automatically - logger.log_message(f"⚠️ WARNING: Missing tables detected in {db_type.upper()} database: {missing_tables}", logging.WARNING) - logger.log_message("🛡️ SAFETY: Not creating tables automatically on PostgreSQL/RDS", logging.INFO) - logger.log_message("📋 Please ensure these tables exist in your RDS database:", logging.INFO) - for table in missing_tables: - logger.log_message(f" - {table}", logging.INFO) - - # Continue without failing - the app might still work with existing tables - if 'agent_templates' in missing_tables or 'user_template_preferences' in missing_tables: - logger.log_message("⚠️ Template functionality may not work without agent_templates and user_template_preferences tables", logging.WARNING) - else: - logger.log_message(f"✅ All required tables exist in {db_type.upper()} database", logging.INFO) - -def verify_template_data(): - """Verify that agent templates are populated. Safe for all database types.""" - logger.log_message("📋 Verifying template data...", logging.INFO) - - session = session_factory() - try: - # Check if agent_templates table exists before querying - if not check_table_exists('agent_templates'): - logger.log_message("⚠️ agent_templates table does not exist, skipping template verification", logging.WARNING) - return - - template_count = session.query(AgentTemplate).filter(AgentTemplate.is_active == True).count() - logger.log_message(f"📊 Found {template_count} active templates", logging.INFO) - - if template_count == 0: - logger.log_message("🔧 No templates found, populating...", logging.INFO) - try: - populate_templates() - - # Verify population worked - new_count = session.query(AgentTemplate).filter(AgentTemplate.is_active == True).count() - logger.log_message(f"✅ Templates populated. Total active templates: {new_count}", logging.INFO) - except Exception as e: - logger.log_message(f"❌ Template population failed: {e}", logging.ERROR) - logger.log_message("⚠️ App will continue but template functionality may not work", logging.WARNING) - else: - logger.log_message("✅ Templates already populated", logging.INFO) - - except Exception as e: - logger.log_message(f"❌ Error verifying templates: {e}", logging.ERROR) - logger.log_message("⚠️ Template verification failed, but app will continue", logging.WARNING) - finally: - session.close() - -def test_template_api_functionality(): - """Test that template-related database operations work. Safe for all database types.""" - logger.log_message("🧪 Testing template API functionality...", logging.INFO) - - session = session_factory() - try: - # Check if agent_templates table exists before testing - if not check_table_exists('agent_templates'): - logger.log_message("⚠️ agent_templates table does not exist, skipping API test", logging.WARNING) - return - - # Test basic template query - templates = session.query(AgentTemplate).filter(AgentTemplate.is_active == True).limit(5).all() - logger.log_message(f"✅ Successfully queried {len(templates)} templates", logging.INFO) - - if templates: - sample_template = templates[0] - logger.log_message(f"📄 Sample template: {sample_template.template_name} - {sample_template.display_name}", logging.INFO) - else: - logger.log_message("📭 No templates found in database", logging.INFO) - - except Exception as e: - logger.log_message(f"❌ Template API test failed: {e}", logging.ERROR) - logger.log_message("⚠️ Template API may not work properly", logging.WARNING) - finally: - session.close() - -def run_safe_initialization(): - """Run safe database initialization that respects production databases.""" - db_type = get_database_type() - logger.log_message(f"🚀 Starting SAFE database initialization for {db_type.upper()}...", logging.INFO) - - if db_type == "postgresql": - logger.log_message("🛡️ PostgreSQL/RDS detected - running in SAFE mode", logging.INFO) - logger.log_message("📋 Will only verify schema and populate templates", logging.INFO) - elif db_type == "sqlite": - logger.log_message("💽 SQLite detected - full initialization mode", logging.INFO) - - try: - # Step 1: Initialize database (safe for all types) - logger.log_message("Step 1: Basic database initialization", logging.INFO) - if db_type == "sqlite": - init_db() # Only run full init on SQLite - else: - logger.log_message("Skipping init_db() for PostgreSQL (safety)", logging.INFO) - - # Step 2: Verify schema (safe - only creates tables on SQLite) - logger.log_message("Step 2: Schema verification", logging.INFO) - verify_database_schema() - - # Step 3: Verify template data (safe for all types) - logger.log_message("Step 3: Template data verification", logging.INFO) - verify_template_data() - - # Step 4: Test functionality (safe for all types) - logger.log_message("Step 4: Functionality testing", logging.INFO) - test_template_api_functionality() - - logger.log_message(f"🎉 Safe database initialization completed for {db_type.upper()}!", logging.INFO) - - except Exception as e: - logger.log_message(f"💥 Database initialization failed: {e}", logging.ERROR) - logger.log_message("⚠️ App may still start but some features might not work", logging.WARNING) - # Don't raise - let the app try to start anyway - -if __name__ == "__main__": - run_safe_initialization() \ No newline at end of file diff --git a/scripts/populate_agent_templates.py b/scripts/populate_agent_templates.py deleted file mode 100644 index 95145cc86caa20348bbf350944c543886726a017..0000000000000000000000000000000000000000 --- a/scripts/populate_agent_templates.py +++ /dev/null @@ -1,508 +0,0 @@ -#!/usr/bin/env python3 -""" -SQLite Agent Template Management Script -Similar to manage_templates.py but optimized for local SQLite development. -Reads agents from agents_config.json and manages SQLite database. -""" - -import sys -import os -import json -import requests -from datetime import datetime, UTC -from pathlib import Path - -# Add the project root to the Python path -script_dir = os.path.dirname(os.path.abspath(__file__)) -backend_dir = os.path.dirname(script_dir) -project_root = os.path.dirname(os.path.dirname(backend_dir)) - -# Change to backend directory to ensure proper path resolution -os.chdir(backend_dir) -sys.path.append(backend_dir) - -from src.db.init_db import session_factory, DATABASE_URL -from src.db.schemas.models import AgentTemplate -from sqlalchemy.exc import IntegrityError - -def get_database_type(): - """Detect database type from DATABASE_URL""" - if DATABASE_URL.startswith('postgresql'): - return "postgresql" - elif DATABASE_URL.startswith('sqlite'): - return "sqlite" - else: - return "unknown" - -def load_agents_config(): - """Load agents configuration from agents_config.json""" - # Try multiple possible locations for agents_config.json - possible_paths = [ - os.path.join(backend_dir, 'agents_config.json'), # Backend directory (copied file) - os.path.join(project_root, 'agents_config.json'), # Project root - '/app/agents_config.json', # Container root (HF Spaces) - 'agents_config.json' # Current directory - ] - - config_path = None - for path in possible_paths: - if os.path.exists(path): - config_path = path - print(f"📖 Found agents_config.json at: {config_path}") - break - - if not config_path: - paths_str = '\n '.join(possible_paths) - raise FileNotFoundError(f"agents_config.json not found in any of these locations:\n {paths_str}") - - with open(config_path, 'r', encoding='utf-8') as f: - config = json.load(f) - - return config.get('templates', []) - -def download_icon(icon_url, template_name): - """Download icon from URL and save to frontend directory""" - if not icon_url or not icon_url.startswith('http'): - print(f"⏭️ Skipping icon download for {template_name} (not a URL: {icon_url})") - return icon_url - - try: - # Determine frontend directory - frontend_dir = os.path.join(project_root, 'Auto-Analyst-CS', 'auto-analyst-frontend') - public_dir = os.path.join(frontend_dir, 'public') - - if not os.path.exists(public_dir): - print(f"⚠️ Frontend public directory not found: {public_dir}") - return icon_url - - # Parse the path from icon_url - if '/icons/templates/' in icon_url: - relative_path = icon_url.split('/icons/templates/')[-1] - icon_dir = os.path.join(public_dir, 'icons', 'templates') - else: - # Fallback: use filename from URL - filename = icon_url.split('/')[-1] - if not filename.endswith(('.svg', '.png', '.jpg', '.jpeg')): - filename += '.svg' - relative_path = filename - icon_dir = os.path.join(public_dir, 'icons', 'templates') - - # Create icon directory if it doesn't exist - os.makedirs(icon_dir, exist_ok=True) - - # Download and save icon - icon_path = os.path.join(icon_dir, relative_path) - - # Skip if already exists - if os.path.exists(icon_path): - print(f"📁 Icon already exists: {relative_path}") - return f"/icons/templates/{relative_path}" - - response = requests.get(icon_url, timeout=10) - response.raise_for_status() - - with open(icon_path, 'wb') as f: - f.write(response.content) - - print(f"📥 Downloaded icon: {relative_path}") - return f"/icons/templates/{relative_path}" - - except Exception as e: - print(f"❌ Failed to download icon for {template_name}: {str(e)}") - return icon_url - -def sync_agents_from_config(): - """Synchronize agents from agents_config.json to SQLite database""" - session = session_factory() - db_type = get_database_type() - - # if db_type != "sqlite": - # print(f"⚠️ This script is designed for SQLite, but detected {db_type}") - # print("Consider using manage_templates.py for PostgreSQL") - # return - - try: - # Load configuration - print(f"📖 Loading agents from agents_config.json...") - templates_config = load_agents_config() - - if not templates_config: - print("❌ No templates found in agents_config.json") - return - - # Track statistics - created_count = 0 - updated_count = 0 - skipped_count = 0 - - print(f"🔍 Processing {len(templates_config)} templates for SQLite database") - print(f"📋 Database URL: {DATABASE_URL}") - - # Group templates by category for display - categories = {} - for template_data in templates_config: - category = template_data.get('category', 'Uncategorized') - if category not in categories: - categories[category] = [] - categories[category].append(template_data) - - # Process templates by category - for category, templates in categories.items(): - print(f"\n📁 {category}:") - - for template_data in templates: - template_name = template_data["template_name"] - - # Check if template already exists - existing = session.query(AgentTemplate).filter( - AgentTemplate.template_name == template_name - ).first() - - # Download icon if it's a URL - icon_url = template_data.get("icon_url", "") - if icon_url.startswith('http'): - icon_url = download_icon(icon_url, template_name) - - if existing: - # Update existing template - existing.display_name = template_data["display_name"] - existing.description = template_data["description"] - existing.icon_url = icon_url - existing.prompt_template = template_data["prompt_template"] - existing.category = template_data.get("category", "Uncategorized") - existing.is_premium_only = template_data.get("is_premium_only", False) - existing.is_active = template_data.get("is_active", True) - existing.variant_type = template_data.get("variant_type", "individual") - existing.base_agent = template_data.get("base_agent", template_name) - existing.updated_at = datetime.now(UTC) - - variant_icon = "🤖" if template_data.get("variant_type") == "planner" else "👤" - premium_icon = "🔒" if template_data.get("is_premium_only") else "🆓" - print(f"🔄 Updated: {template_name} {variant_icon} {premium_icon}") - updated_count += 1 - else: - # Create new template - template = AgentTemplate( - template_name=template_name, - display_name=template_data["display_name"], - description=template_data["description"], - icon_url=icon_url, - prompt_template=template_data["prompt_template"], - category=template_data.get("category", "Uncategorized"), - is_premium_only=template_data.get("is_premium_only", False), - is_active=template_data.get("is_active", True), - variant_type=template_data.get("variant_type", "individual"), - base_agent=template_data.get("base_agent", template_name), - created_at=datetime.now(UTC), - updated_at=datetime.now(UTC) - ) - - session.add(template) - variant_icon = "🤖" if template_data.get("variant_type") == "planner" else "👤" - premium_icon = "🔒" if template_data.get("is_premium_only") else "🆓" - print(f"✅ Created: {template_name} {variant_icon} {premium_icon}") - created_count += 1 - - # Handle removals if specified in config - remove_list = [] - # Re-load the full config to check for removals - try: - full_config_path = None - possible_paths = [ - os.path.join(backend_dir, 'agents_config.json'), - os.path.join(project_root, 'agents_config.json'), - '/app/agents_config.json', - 'agents_config.json' - ] - - for path in possible_paths: - if os.path.exists(path): - full_config_path = path - break - - if full_config_path: - with open(full_config_path, 'r', encoding='utf-8') as f: - full_config = json.load(f) - if 'remove' in full_config: - remove_list = full_config['remove'] - except Exception as e: - print(f"⚠️ Could not load removal list: {e}") - - # Remove templates marked for removal - if remove_list: - print(f"\n🗑️ --- Processing Removals ---") - for template_name in remove_list: - existing = session.query(AgentTemplate).filter( - AgentTemplate.template_name == template_name - ).first() - - if existing: - session.delete(existing) - print(f"🗑️ Removed: {template_name}") - else: - print(f"⏭️ Skipping removal: {template_name} (not found)") - - # Commit all changes - session.commit() - - print(f"\n📊 --- Summary ---") - print(f"✅ Templates created: {created_count}") - print(f"🔄 Templates updated: {updated_count}") - print(f"⏭️ Templates skipped: {skipped_count}") - - # Show total count in database - total_count = session.query(AgentTemplate).count() - free_count = session.query(AgentTemplate).filter(AgentTemplate.is_premium_only == False).count() - premium_count = session.query(AgentTemplate).filter(AgentTemplate.is_premium_only == True).count() - individual_count = session.query(AgentTemplate).filter(AgentTemplate.variant_type == 'individual').count() - planner_count = session.query(AgentTemplate).filter(AgentTemplate.variant_type == 'planner').count() - - print(f"🗄️ Total templates in database: {total_count}") - print(f"🆓 Free templates: {free_count}") - print(f"🔒 Premium templates: {premium_count}") - print(f"👤 Individual variants: {individual_count}") - print(f"🤖 Planner variants: {planner_count}") - - except Exception as e: - session.rollback() - print(f"❌ Error syncing templates: {str(e)}") - raise - finally: - session.close() - -def list_templates(): - """List all existing templates in the database""" - session = session_factory() - - try: - templates = session.query(AgentTemplate).order_by( - AgentTemplate.category, - AgentTemplate.is_premium_only, - AgentTemplate.template_name - ).all() - - if not templates: - print("No templates found in database.") - return - - print(f"\n--- Existing Templates ({len(templates)} total) ---") - - current_category = None - for template in templates: - if template.category != current_category: - current_category = template.category - print(f"\n📁 {current_category}:") - - status = "🔒 Premium" if template.is_premium_only else "🆓 Free" - active = "✅ Active" if template.is_active else "❌ Inactive" - variant = getattr(template, 'variant_type', 'individual') - variant_icon = "🤖" if variant == "planner" else "👤" - - print(f" • {template.template_name} ({template.display_name})") - print(f" {status} - {active} - {variant_icon} {variant}") - print(f" 📝 {template.description}") - - except Exception as e: - print(f"❌ Error listing templates: {str(e)}") - finally: - session.close() - -def remove_all_templates(): - """Remove all templates from database (for testing)""" - session = session_factory() - - try: - deleted_count = session.query(AgentTemplate).delete() - session.commit() - print(f"🗑️ Removed {deleted_count} templates from database") - - except Exception as e: - session.rollback() - print(f"❌ Error removing templates: {str(e)}") - finally: - session.close() - -def validate_config(): - """Validate the agents_config.json structure""" - try: - templates_config = load_agents_config() - - print(f"📋 Validating agents_config.json...") - print(f"✅ Found {len(templates_config)} templates") - - # Check required fields - required_fields = ['template_name', 'display_name', 'description', 'prompt_template'] - issues = [] - - for i, template in enumerate(templates_config): - for field in required_fields: - if field not in template: - issues.append(f"Template {i}: Missing required field '{field}'") - - if issues: - print(f"❌ Validation issues found:") - for issue in issues: - print(f" • {issue}") - else: - print(f"✅ Configuration is valid") - - # Show summary by category - categories = {} - for template in templates_config: - category = template.get('category', 'Uncategorized') - if category not in categories: - categories[category] = {'free': 0, 'premium': 0, 'individual': 0, 'planner': 0} - - if template.get('is_premium_only', False): - categories[category]['premium'] += 1 - else: - categories[category]['free'] += 1 - - if template.get('variant_type', 'individual') == 'planner': - categories[category]['planner'] += 1 - else: - categories[category]['individual'] += 1 - - print(f"\n📊 Summary by category:") - for category, counts in categories.items(): - total = counts['free'] + counts['premium'] - print(f" 📁 {category}: {total} templates") - print(f" 🆓 Free: {counts['free']} | 🔒 Premium: {counts['premium']}") - print(f" 👤 Individual: {counts['individual']} | 🤖 Planner: {counts['planner']}") - - except Exception as e: - print(f"❌ Error validating config: {str(e)}") - -def create_minimal_templates(): - """Create a minimal set of essential templates for container environments""" - session = session_factory() - - try: - print("🔧 Creating minimal template set...") - - # Define minimal essential templates - minimal_templates = [ - { - "template_name": "preprocessing_agent", - "display_name": "Data Preprocessing Agent", - "description": "Cleans and prepares DataFrame using Pandas and NumPy", - "icon_url": "/icons/templates/preprocessing_agent.svg", - "category": "Data Manipulation", - "is_premium_only": False, - "variant_type": "individual", - "base_agent": "preprocessing_agent", - "is_active": True, - "prompt_template": "You are a preprocessing agent that cleans and prepares data using Pandas and NumPy. Handle missing values, detect column types, and convert date strings to datetime. Generate clean Python code for data preprocessing based on the user's analysis goals." - }, - { - "template_name": "data_viz_agent", - "display_name": "Data Visualization Agent", - "description": "Creates interactive visualizations using Plotly", - "icon_url": "/icons/templates/data_viz_agent.svg", - "category": "Data Visualization", - "is_premium_only": False, - "variant_type": "individual", - "base_agent": "data_viz_agent", - "is_active": True, - "prompt_template": "You are a data visualization agent. Create interactive visualizations using Plotly based on user requirements. Generate appropriate chart types, apply styling, and ensure visualizations effectively communicate insights." - }, - { - "template_name": "sk_learn_agent", - "display_name": "Machine Learning Agent", - "description": "Trains ML models using scikit-learn", - "icon_url": "/icons/templates/sk_learn_agent.svg", - "category": "Data Modelling", - "is_premium_only": False, - "variant_type": "individual", - "base_agent": "sk_learn_agent", - "is_active": True, - "prompt_template": "You are a machine learning agent. Use scikit-learn to train and evaluate ML models including classification, regression, and clustering. Provide feature importance insights and model performance metrics." - } - ] - - created_count = 0 - - for template_data in minimal_templates: - template_name = template_data["template_name"] - - # Check if template already exists - existing = session.query(AgentTemplate).filter( - AgentTemplate.template_name == template_name - ).first() - - if not existing: - template = AgentTemplate( - template_name=template_name, - display_name=template_data["display_name"], - description=template_data["description"], - icon_url=template_data["icon_url"], - prompt_template=template_data["prompt_template"], - category=template_data["category"], - is_premium_only=template_data["is_premium_only"], - is_active=template_data["is_active"], - variant_type=template_data["variant_type"], - base_agent=template_data["base_agent"], - created_at=datetime.now(UTC), - updated_at=datetime.now(UTC) - ) - - session.add(template) - print(f"✅ Created minimal template: {template_name}") - created_count += 1 - else: - print(f"⏭️ Template already exists: {template_name}") - - session.commit() - print(f"📊 Created {created_count} minimal templates") - - except Exception as e: - session.rollback() - print(f"❌ Error creating minimal templates: {str(e)}") - raise - finally: - session.close() - -def populate_templates(): - """Legacy compatibility function for backward compatibility""" - print("⚠️ Legacy populate_templates() called - checking for agents_config.json...") - - # Check if agents_config.json exists anywhere - possible_paths = [ - os.path.join(backend_dir, 'agents_config.json'), - os.path.join(project_root, 'agents_config.json'), - '/app/agents_config.json', - 'agents_config.json' - ] - - config_exists = any(os.path.exists(path) for path in possible_paths) - - if config_exists: - print("📖 Found agents_config.json - using sync_agents_from_config()") - sync_agents_from_config() - else: - print("⚠️ agents_config.json not found - using fallback minimal templates") - print("💡 Creating essential templates for container environment") - create_minimal_templates() - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="SQLite Agent Template Management") - parser.add_argument("action", choices=["sync", "list", "remove-all", "validate"], - help="Action to perform") - - args = parser.parse_args() - - if args.action == "sync": - print("🚀 Synchronizing agents from agents_config.json to SQLite...") - sync_agents_from_config() - elif args.action == "list": - list_templates() - elif args.action == "validate": - validate_config() - elif args.action == "remove-all": - confirm = input("⚠️ Are you sure you want to remove ALL templates? (yes/no): ") - if confirm.lower() == "yes": - remove_all_templates() - else: - print("Operation cancelled.") \ No newline at end of file diff --git a/scripts/setup_analytics_data.py b/scripts/setup_analytics_data.py new file mode 100644 index 0000000000000000000000000000000000000000..a93a66d5e501c22e27dbbb49b1785b02b41aa654 --- /dev/null +++ b/scripts/setup_analytics_data.py @@ -0,0 +1,114 @@ +import sys +import os +from datetime import datetime, timedelta +import sqlite3 + +# Add parent directory to path so we can import modules +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.db.schemas.models import ModelUsage, User +from src.db.init_db import session_factory, init_db +from scripts.generate_test_data import generate_test_data +from scripts.create_test_user import create_test_users + +def check_database(): + """Check if database exists and has the required tables""" + db_path = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "chat_database.db" + ) + + if not os.path.exists(db_path): + print(f"Database file not found at {db_path}") + print("Creating database...") + init_db() + return False + + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + + # Check if model_usage table exists + cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='model_usage'") + if not cursor.fetchone(): + print("model_usage table does not exist.") + print("Creating tables...") + init_db() + return False + + # Count records in model_usage table + cursor.execute("SELECT COUNT(*) FROM model_usage") + count = cursor.fetchone()[0] + print(f"Found {count} records in model_usage table") + + conn.close() + return count > 0 + +def setup_analytics(): + """Set up analytics data for testing""" + # Check database status + has_data = check_database() + + # Create test users + print("\nCreating test users...") + create_test_users() + + # Generate model usage data if needed + if not has_data: + print("\nGenerating model usage test data...") + generate_test_data(100) + else: + print("\nDatabase already has model usage data. Skipping generation.") + choice = input("Generate additional data anyway? (y/n): ") + if choice.lower() == 'y': + records = int(input("How many records to generate? [default: 100]: ") or "100") + generate_test_data(records) + + # Verify data was created + session = session_factory() + try: + # Count model usage records + usage_count = session.query(ModelUsage).count() + print(f"\nTotal model usage records: {usage_count}") + + # Count users + user_count = session.query(User).count() + print(f"Total users: {user_count}") + + # Check recent data + yesterday = datetime.utcnow() - timedelta(days=1) + recent_count = session.query(ModelUsage).filter(ModelUsage.timestamp >= yesterday).count() + print(f"Records from the last 24 hours: {recent_count}") + + # Get model breakdown + models = {} + providers = {} + + for usage in session.query(ModelUsage).all(): + if usage.model_name not in models: + models[usage.model_name] = 0 + models[usage.model_name] += 1 + + if usage.provider not in providers: + providers[usage.provider] = 0 + providers[usage.provider] += 1 + + print("\nModel breakdown:") + for model, count in models.items(): + print(f" {model}: {count} records") + + print("\nProvider breakdown:") + for provider, count in providers.items(): + print(f" {provider}: {count} records") + + finally: + session.close() + + print("\nSetup complete!") + print("To access the analytics dashboard:") + print("1. Make sure the backend server is running") + print("2. In your browser, set localStorage.adminApiKey to 'default-admin-key-change-me'") + print(" (or the value in your ADMIN_API_KEY environment variable)") + print("3. Go to http://localhost:3000/analytics/dashboard") + +if __name__ == "__main__": + setup_analytics() \ No newline at end of file diff --git a/scripts/test_agent_parsing.py b/scripts/test_agent_parsing.py new file mode 100644 index 0000000000000000000000000000000000000000..a18c30c7d5a0f202881a8ad6f4903f247a9c8063 --- /dev/null +++ b/scripts/test_agent_parsing.py @@ -0,0 +1,42 @@ + +def parse_agents(agent_string): + """ + Parse a string containing agent names separated by ->, (, ), or commas + and return a list of agent names. + """ + if not agent_string or not agent_string.strip(): + return [] + + # Replace parentheses with spaces to handle cases with parentheses + import re + cleaned_string = re.sub(r'\(.*?\)', '', agent_string) + + # Split by -> to get individual agent segments + agent_segments = cleaned_string.split('->') + + # Process each segment to extract agent names + agents = [] + for segment in agent_segments: + # Split by comma and strip whitespace + segment_agents = [agent.strip() for agent in segment.split(',') if agent.strip()] + agents.extend(segment_agents) + + return agents + +sample = "preprocessing_agent -> statistical_analytics_agent" +agents = parse_agents(sample) +print(agents) + +# Test with different formats +test_samples = [ + "preprocessing_agent -> data_viz_agent", + "preprocessing_agent(dataset, goal)", + "preprocessing_agent -> statistical_analytics_agent, data_viz_agent", + "preprocessing_agent, statistical_analytics_agent -> data_viz_agent", + "", + "preprocessing_agent(dataset, goal) -> data_viz_agent(dataset)" +] + +for test in test_samples: + print(f"Input: {test}") + print(f"Output: {parse_agents(test)}") \ No newline at end of file diff --git a/scripts/test_model_usage.py b/scripts/test_model_usage.py new file mode 100644 index 0000000000000000000000000000000000000000..73b104786afe51929de10e6ddad9d0fbf370b971 --- /dev/null +++ b/scripts/test_model_usage.py @@ -0,0 +1,114 @@ +import requests +import time +import uuid +import os +import sys + +# Add the parent directory to the path so we can import modules +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from ai_manager import AI_Manager +from init_db import session_factory, ModelUsage + +def test_ai_manager_directly(): + """Test the AI Manager directly to see if it records usage""" + print("Starting direct AI Manager test...") + + # Create a test AI Manager + ai_manager = AI_Manager() + + # Generate a unique user ID and chat ID for testing + test_user_id = int(time.time()) % 10000 + test_chat_id = int(uuid.uuid4().int % 100000) + test_model = "test-direct-model" + + print(f"Using test user_id: {test_user_id}, chat_id: {test_chat_id}") + + # Call save_usage_to_db directly + ai_manager.save_usage_to_db( + user_id=test_user_id, + chat_id=test_chat_id, + model_name=test_model, + provider="Test Provider", + prompt_tokens=150, + completion_tokens=250, + total_tokens=400, + query_size=600, + response_size=1200, + cost=0.002, + request_time_ms=200, + is_streaming=False + ) + + print("Usage data saved directly via AI Manager") + + # Check if it was actually saved + session = session_factory() + try: + # Count records for our test user + count = session.query(ModelUsage).filter( + ModelUsage.user_id == test_user_id, + ModelUsage.model_name == test_model + ).count() + + if count > 0: + print(f"SUCCESS: Found {count} records for test user {test_user_id}") + records = session.query(ModelUsage).filter( + ModelUsage.user_id == test_user_id + ).all() + + for record in records: + print(f" - {record.usage_id}: {record.model_name}, {record.total_tokens} tokens, ${record.cost}") + else: + print(f"FAILURE: No records found for test user {test_user_id}") + print("Check the database connection and ModelUsage table") + finally: + session.close() + +def test_api_endpoint(): + """Test the API endpoint that creates test usage records""" + print("\nTesting API endpoint for creating test usage...") + + admin_key = os.getenv("ADMIN_API_KEY", "admin-api-key-change-me") + base_url = "http://localhost:8000" + + # Generate a unique user ID for testing + test_user_id = int(time.time()) % 10000 + test_model = "test-api-model" + + # Call the API endpoint + response = requests.post( + f"{base_url}/analytics/debug/create_test_usage", + params={ + "user_id": test_user_id, + "model_name": test_model + }, + headers={ + "X-Admin-API-Key": admin_key + } + ) + + if response.status_code == 200: + data = response.json() + print(f"SUCCESS: Created test usage via API endpoint") + print(f"Response: {data}") + + # Verify the record was created + verification = requests.get( + f"{base_url}/analytics/debug/user_usage/{test_user_id}", + headers={ + "X-Admin-API-Key": admin_key + } + ).json() + + print(f"User {test_user_id} has {verification.get('total_records', 0)} records") + if verification.get('recent_records'): + print(f"Recent records: {verification['recent_records']}") + else: + print(f"FAILURE: API endpoint returned status {response.status_code}") + print(f"Response: {response.text}") + +if __name__ == "__main__": + # Run both tests + test_ai_manager_directly() + test_api_endpoint() \ No newline at end of file diff --git a/scripts/test_token_counting.py b/scripts/test_token_counting.py new file mode 100644 index 0000000000000000000000000000000000000000..133ea1dc6f654fd6e1c96b22cc1635d8117e43e3 --- /dev/null +++ b/scripts/test_token_counting.py @@ -0,0 +1,176 @@ +import sys +from pathlib import Path +sys.path.append(str(Path(__file__).resolve().parent.parent)) + +from src.managers.ai_manager import AI_Manager + +ai_manager = AI_Manager() + +input_prompt = "What are the key insights from this data? " + +output = """ +## Analytical Planner + +### Reasoning +To extract key insights from the housing dataset, it is essential to first preprocess the data to ensure its quality and suitability for analysis. Following preprocessing, performing statistical analysis will help identify significant patterns and relationships within the data. Finally, visualizing the results will provide a clear and interpretable representation of the insights derived. + + +## Preprocessing Agent + +### Reasoning +To derive key insights from the Housing Dataset, it is essential to start with data preprocessing and introductory analysis. This involves categorizing the data into numeric and categorical features, handling missing values appropriately, and preparing the data for further exploration. By organizing the data effectively and understanding the relationships between different variables, we can uncover meaningful patterns and trends that inform our analysis objectives. + +### Code Implementation +```python +import numpy as np +import pandas as pd + +# Define a safe datetime conversion function +def safe_to_datetime(date): + try: + return pd.to_datetime(date, errors='coerce', cache=False) + except (ValueError, TypeError): + return pd.NaT + +# Create copies of the original dataframe +df_copy = df.copy() + +# Identify categorical and numeric columns +categorical_columns = df_copy.select_dtypes(include=['object', 'category']).columns.tolist() +numeric_columns = df_copy.select_dtypes(include=[np.number]).columns.tolist() + +# Handle missing values +# For numeric columns, fill missing values with the median +df_copy[numeric_columns] = df_copy[numeric_columns].fillna(df_copy[numeric_columns].median()) + +# For categorical columns, fill missing values with the mode +for col in categorical_columns: + df_copy[col] = df_copy[col].fillna(df_copy[col].mode()[0]) + +# Convert string dates to datetime format if any date columns exist +# Example: If there's a column named 'date', uncomment the following lines +# if 'date' in df_copy.columns: +# df_copy['date'] = df_copy['date'].apply(safe_to_datetime) + +# Create a correlation matrix for numeric columns +correlation_matrix = df_copy[numeric_columns].corr() + +# Display the first few rows of the processed dataframe and the correlation matrix +print(df_copy.head()) +print("\nCorrelation Matrix:\n", correlation_matrix) +``` + +### Summary +1. **Import Libraries**: The necessary libraries, NumPy and Pandas, are imported to handle numerical operations and data manipulation. + +2. **Datetime Conversion Function**: A `safe_to_datetime` function is defined to gracefully handle the conversion of string dates to datetime objects, ensuring that any errors during conversion result in `NaT` (Not a Time) values without interrupting the workflow. + +3. **Dataframe Copy**: A copy of the original dataframe `df` is created to preserve the original data during preprocessing. + +4. **Identify Column Types**: + - **Categorical Columns**: Identified using `select_dtypes` to include data types `object` and `category`. + - **Numeric Columns**: Identified by selecting columns with numeric data types using `np.number`. + +5. **Handle Missing Values**: + - **Numeric Columns**: Missing values are filled with the median of each numeric column to maintain the central tendency without being skewed by outliers. + - **Categorical Columns**: Missing values are filled with the mode (most frequent value) to preserve the most common category. + +6. **Datetime Conversion**: Although the current dataset does not contain date columns, the preparation is made to handle such conversions if needed in the future by checking for date-specific columns and applying the `safe_to_datetime` function. + +7. **Correlation Matrix**: A correlation matrix is generated for the numeric columns to identify the strength and direction of relationships between different numerical variables, which is crucial for understanding how features interact with each other. + +8. **Display Outputs**: The first few rows of the processed dataframe and the correlation matrix are printed to provide an initial view of the data and its interrelationships, laying the groundwork for deeper exploratory data analysis. + + +## Statistical Analytics Agent + +### Reasoning +To extract key insights from the Housing Dataset, we will perform a multiple linear regression analysis. This will allow us to identify which factors significantly influence housing prices. By analyzing both numerical variables (such as area, bedrooms, bathrooms) and categorical variables (like mainroad, furnishingstatus), we can determine the key drivers of price variations. Proper handling of missing values and categorical data will ensure the robustness of our analysis. + +### Code Implementation +```python +import pandas as pd +import numpy as np +import statsmodels.api as sm + +# Create a copy of the dataframe +df_copy = df.copy() + +try: + # Define the dependent variable + y = df_copy['price'] + + # Define the independent variables + X = df_copy.drop('price', axis=1) + + # Check for missing values and drop rows with any missing values + X = X.dropna() + y = y.loc[X.index] + + # Convert categorical variables to 'category' dtype + categorical_columns = X.select_dtypes(include=['object', 'category']).columns + for col in categorical_columns: + X[col] = X[col].astype('category') + + # Add a constant term to the predictors + X = sm.add_constant(X) + + # Fit the OLS regression model + model = sm.OLS(y.astype(float), X.astype(float)).fit() + + # Print the model summary + print(model.summary()) + +except Exception as e: + print(f"An error occurred: {e}") +``` + +### Summary +The code performs a multiple linear regression analysis to identify the key factors affecting housing prices. It begins by creating a copy of the original dataframe to preserve the data. The dependent variable `price` is separated from the independent variables. Rows with any missing values are removed to ensure data quality. Categorical variables are converted to the 'category' data type, which allows statsmodels to handle them appropriately in the regression. A constant term is added to the independent variables to account for the intercept in the model. The Ordinary Least Squares (OLS) regression model is then fitted using the prepared data. Finally, the summary of the regression model is printed, providing detailed statistics on the significance and impact of each predictor on the housing price. + + +## Data Viz Agent + +### Reasoning +To extract key insights from the housing dataset, it's important to analyze the relationship between house area and price, as well as understand how factors like the number of bedrooms and bathrooms influence pricing. A scatter plot of area versus price can reveal trends and correlations, while color-coding by the number of bedrooms can provide additional depth to the analysis. + +### Code Implementation +```python +import plotly.express as px + +# Check if the dataset has more than 50,000 rows and sample if necessary +if len(df) > 50000: + df_sample = df.sample(n=50000) +else: + df_sample = df.copy() + +# Create a scatter plot of area vs price, colored by number of bedrooms +fig = px.scatter(df_sample, + x='area', + y='price', + color='bedrooms', + title='Housing Price vs Area Colored by Number of Bedrooms', + labels={'area': 'Area (sq ft)', 'price': 'Price'}, + hover_data=['bathrooms', 'stories']) + +# Update layout for clarity and styling +fig.update_layout( + template='plotly_white', + xaxis_title='Area (sq ft)', + yaxis_title='Price', + legend_title='Bedrooms' +) + +fig.to_html(full_html=False) +``` + +### Summary +The scatter plot visualizes the relationship between the area of the houses and their prices. By color-coding the points based on the number of bedrooms, it becomes easier to observe how bedroom count correlates with both area and price. This helps in identifying trends such as whether larger houses with more bedrooms tend to be priced higher. +""" + + +input_tokens = len(ai_manager.tokenizer.encode(input_prompt)) +output_tokens = len(ai_manager.tokenizer.encode(output)) + +print(f"Input tokens: {input_tokens}") +print(f"Output tokens: {output_tokens}") \ No newline at end of file diff --git a/scripts/tier_maker.py b/scripts/tier_maker.py index 0d0dd547b3ade195bc9ed54745546008d1ff3eb8..fb1670ac0149caf40f7a0f499962a0abe2c6c079 100644 --- a/scripts/tier_maker.py +++ b/scripts/tier_maker.py @@ -2,16 +2,8 @@ from src.utils.model_registry import MODEL_COSTS, MODEL_TIERS # divide models in 3 tiers based on cost per 1k tokens # tier 1: < $0.0005 -# tier 2: < $0.001 -# tier 3: > $0.05 -# tier 4: > $0.1 - -TIERS_COST = { - "tier1": 0.0005, - "tier2": 0.001, - "tier3": 0.05, - "tier4": 0.1 -} +# tier 2: $0.0005 - $0.001 +# tier 3: > $0.001 def get_tier(model_name): for provider, models in MODEL_COSTS.items(): @@ -24,7 +16,7 @@ def get_tier_1(): tier_1 = [] for provider, models in MODEL_COSTS.items(): for model, cost in models.items(): - if cost["input"] + cost["output"] < TIERS_COST["tier1"]: + if cost["input"] + cost["output"] < 0.0005: tier_1.append(model) return tier_1 @@ -32,7 +24,7 @@ def get_tier_2(): tier_2 = [] for provider, models in MODEL_COSTS.items(): for model, cost in models.items(): - if cost["input"] + cost["output"] >= TIERS_COST["tier1"] and cost["input"] + cost["output"] < TIERS_COST["tier2"]: + if cost["input"] + cost["output"] >= 0.0005 and cost["input"] + cost["output"] < 0.001: tier_2.append(model) return tier_2 @@ -40,17 +32,9 @@ def get_tier_3(): tier_3 = [] for provider, models in MODEL_COSTS.items(): for model, cost in models.items(): - if cost["input"] + cost["output"] >= TIERS_COST["tier2"] and cost["input"] + cost["output"] < TIERS_COST["tier3"]: + if cost["input"] + cost["output"] >= 0.001: tier_3.append(model) - return tier_3 - -def get_tier_4(): - tier_4 = [] - for provider, models in MODEL_COSTS.items(): - for model, cost in models.items(): - if cost["input"] + cost["output"] >= TIERS_COST["tier3"]: - tier_4.append(model) - return tier_4 + return tier_3 # Print current tier definitions from registry import json @@ -74,11 +58,6 @@ model_tiers = { "name": "Premium", "credits": 5, "models": get_tier_3() - }, - "tier4": { - "name": "Premium Plus", - "credits": 10, - "models": get_tier_4() } } diff --git a/scripts/verify_session_state.py b/scripts/verify_session_state.py new file mode 100644 index 0000000000000000000000000000000000000000..b25d3e98597101780dce3527a5c1faaf3c8e4d4c --- /dev/null +++ b/scripts/verify_session_state.py @@ -0,0 +1,76 @@ +import requests +import uuid +import time + +def test_session_workflow(): + # Base URL + base_url = "http://localhost:8000" + + # Create a unique session ID + session_id = str(uuid.uuid4()) + print(f"Testing with session ID: {session_id}") + + # Step 1: Create a user and associate with session + print("\n1. Creating user...") + username = f"test_user_{session_id[:8]}" + email = f"{username}@example.com" + login_response = requests.post( + f"{base_url}/auth/login", + json={"username": username, "email": email, "session_id": session_id} + ).json() + print(f"Login response: {login_response}") + + # Step 2: Verify session state + print("\n2. Verifying session state...") + session_state = requests.get(f"{base_url}/debug/session/{session_id}").json() + print(f"Session state: {session_state}") + + # Step 3: Make a chat request with this session + print("\n3. Making a test chat request...") + chat_response = requests.post( + f"{base_url}/chat/data_explorer", + json={"query": "Show me a summary of the dataset"}, + params={"session_id": session_id} + ).json() + print(f"Chat response received: {len(str(chat_response))} bytes") + + # Step 4: Verify session state again + print("\n4. Verifying session state after chat...") + session_state = requests.get(f"{base_url}/debug/session/{session_id}").json() + print(f"Updated session state: {session_state}") + + # Step 5: Check analytics data + print("\n5. Checking analytics data...") + time.sleep(1) # Wait a moment for data to be saved + admin_key = "default-admin-key-change-me" # Adjust to your admin key + + # Check general model usage + analytics_response = requests.get( + f"{base_url}/analytics/debug/model_usage", + headers={"X-Admin-API-Key": admin_key} + ).json() + print(f"Total records in database: {analytics_response.get('total_records', 0)}") + + if 'sample_records' in analytics_response and analytics_response['sample_records']: + print(f"Latest usage records: {analytics_response['sample_records']}") + else: + print("No sample records found in general model usage") + + # Check user-specific usage + user_id = login_response.get('user_id') + if user_id: + user_usage = requests.get( + f"{base_url}/analytics/debug/user_usage/{user_id}", + headers={"X-Admin-API-Key": admin_key} + ).json() + print(f"User {user_id} has {user_usage.get('total_records', 0)} usage records") + + if user_usage.get('recent_records'): + print(f"Recent usage: {user_usage['recent_records']}") + else: + print(f"No usage records found for user {user_id}") + + print("\nTest completed!") + +if __name__ == "__main__": + test_session_workflow() \ No newline at end of file diff --git a/src/agents/agents.py b/src/agents/agents.py index 2c21c3f198c9d32d384ee8e508ec2f784e4dea87..3746616f007dc8c0e321c5d271b0ce71e6b7ce04 100644 --- a/src/agents/agents.py +++ b/src/agents/agents.py @@ -1,386 +1,53 @@ import dspy import src.agents.memory_agents as m import asyncio +from concurrent.futures import ThreadPoolExecutor +import os from dotenv import load_dotenv import logging from src.utils.logger import Logger -from src.utils.model_registry import small_lm, mid_lm -import json - load_dotenv() -logger = Logger("agents", see_time=True, console_log=True) - - - -# === CUSTOM AGENT FUNCTIONALITY === -def create_custom_agent_signature(agent_name, description, prompt_template, category=None): - """ - Dynamically creates a dspy.Signature class for custom agents. - Has to be tested - - Args: - agent_name: Name of the custom agent (e.g., 'pytorch_agent') - description: Short description for agent selection - prompt_template: Main prompt/instructions for agent behavior - category: Agent category from database (e.g., 'Visualization', 'Modelling', 'Data Manipulation') - - Returns: - A dspy.Signature class with the custom prompt and standard input/output fields - """ - - # Check if this is a visualization agent to determine input fields - # First check category, then fallback to name-based detection - if category and category.lower() == 'visualization': - is_viz_agent = True - else: - is_viz_agent = 'viz' in agent_name.lower() or 'visual' in agent_name.lower() or 'plot' in agent_name.lower() or 'chart' in agent_name.lower() - - # Standard input/output fields that match the unified agent signatures - class_attributes = { - '__doc__': prompt_template, # The custom prompt becomes the docstring - 'goal': dspy.InputField(desc="User-defined goal which includes information about data and task they want to perform"), - 'dataset': dspy.InputField(desc="Provides information about the data in the data frame. Only use column names and dataframe_name as in this context"), - 'plan_instructions': dspy.InputField(desc="Agent-level instructions about what to create and receive", default=""), - 'code': dspy.OutputField(desc="Generated Python code for the analysis"), - 'summary': dspy.OutputField(desc="A concise bullet-point summary of what was done and key results") - } - - - # Add styling_index for visualization agents - if is_viz_agent: - class_attributes['styling_index'] = dspy.InputField(desc='Provides instructions on how to style outputs and formatting') - - # Create the dynamic signature class - CustomAgentSignature = type(agent_name, (dspy.Signature,), class_attributes) - return CustomAgentSignature - -def load_user_enabled_templates_from_db(user_id, db_session): - """ - Load template agents that are enabled for a specific user from the database. - Default agents are enabled by default unless explicitly disabled by user preference. - - Args: - user_id: ID of the user - db_session: Database session - - Returns: - Dict of template agent signatures keyed by template name - """ - try: - from src.db.schemas.models import AgentTemplate, UserTemplatePreference - - agent_signatures = {} - - if not user_id: - return agent_signatures - - # Get list of default agent names that should be enabled by default - default_agent_names = [ - "preprocessing_agent", - "statistical_analytics_agent", - "sk_learn_agent", - "data_viz_agent" - ] - - # Get all active templates - all_templates = db_session.query(AgentTemplate).filter( - AgentTemplate.is_active == True - ).all() - - for template in all_templates: - # Check if user has explicitly disabled this template - preference = db_session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - # Determine if template should be enabled by default - is_default_agent = template.template_name in default_agent_names - default_enabled = is_default_agent # Default agents enabled by default, others disabled - - # Template is enabled by default for default agents, disabled for others - is_enabled = preference.is_enabled if preference else default_enabled - - if is_enabled: - # Create dynamic signature for each enabled template - signature = create_custom_agent_signature( - template.template_name, - template.description, - template.prompt_template, - template.category # Pass the category from database - ) - agent_signatures[template.template_name] = signature - - return agent_signatures - - except Exception as e: - logger.log_message(f"Error loading user enabled templates for user {user_id}: {str(e)}", level=logging.ERROR) - return {} - -def load_user_enabled_templates_for_planner_from_db(user_id, db_session): - """ - Load planner variant template agents that are enabled for planner use (max 10, prioritized by usage). - Default planner agents are enabled by default unless explicitly disabled by user preference. - Custom/premium agents require explicit enablement. - - Args: - user_id: ID of the user - db_session: Database session - - Returns: - Dict of template agent signatures keyed by template name (max 10) - """ - # try: - from src.db.schemas.models import AgentTemplate, UserTemplatePreference - from datetime import datetime, UTC - - agent_signatures = {} - - - - # Get list of default planner agent names that should be enabled by default - default_planner_agent_names = [ - "planner_preprocessing_agent", - "planner_statistical_analytics_agent", - "planner_sk_learn_agent", - "planner_data_viz_agent" - ] - # if not user_id: - # return agent_signatures - - # Get all active planner variant templates - all_templates = db_session.query(AgentTemplate).filter( - AgentTemplate.is_active == True, - AgentTemplate.variant_type.in_(['planner', 'both']) - ).all() - - enabled_templates = [] - # Fetch all preferences for the user in a single query for efficiency - preferences = db_session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id - ).all() - preference_map = {pref.template_id: pref for pref in preferences} - - for template in all_templates: - preference = preference_map.get(template.template_id) - is_default_planner_agent = template.template_name in default_planner_agent_names - default_enabled = is_default_planner_agent - is_enabled = preference.is_enabled if preference else default_enabled - - if is_enabled: - enabled_templates.append({ - 'template': template, - 'preference': preference, - 'usage_count': preference.usage_count if preference else 0, - 'last_used_at': preference.last_used_at if preference else None - }) - - # If user has no enabled templates, fall back to default enabled (default planner agents) - if enabled_templates == []: - for template in all_templates: - if template.template_name in default_planner_agent_names: - enabled_templates.append({ - 'template': template, - 'preference': None, - 'usage_count': 0, - 'last_used_at': None - }) - - # Sort by usage (most used first) and limit to 10 - enabled_templates.sort(key=lambda x: (x['usage_count'], x['last_used_at'] or datetime.min.replace(tzinfo=UTC)), reverse=True) - enabled_templates = enabled_templates[:10] - - for item in enabled_templates: - template = item['template'] - # Create dynamic signature for each enabled template - signature = create_custom_agent_signature( - template.template_name, - template.description, - template.prompt_template, - template.category # Pass the category from database - ) - agent_signatures[template.template_name] = signature - - logger.log_message(f"Loaded {len(agent_signatures)} templates for planner", level=logging.DEBUG) - return agent_signatures - - # except Exception as e: - # logger.log_message(f"Error loading planner templates for user {user_id}: {str(e)}", level=logging.ERROR) - # return {} - -def get_all_available_templates(db_session): - """ - Get all available agent templates from the database. - - Args: - db_session: Database session - - Returns: - List of agent template records - """ - try: - from src.db.schemas.models import AgentTemplate - import os - import json - - templates = db_session.query(AgentTemplate).filter( - AgentTemplate.is_active == True - ).all() - - if not templates: - # Try to load from agents_config.json after the main folder (project root or backend dir) - current_dir = os.path.dirname(os.path.abspath(__file__)) - backend_dir = os.path.dirname(current_dir) - project_root = os.path.dirname(backend_dir) - possible_paths = [ - os.path.join(backend_dir, 'agents_config.json'), # backend directory - os.path.join(project_root, 'agents_config.json'), # project root - '/app/agents_config.json' # container root (for Docker/Spaces) - ] - config_path = None - for path in possible_paths: - if os.path.exists(path): - config_path = path - break - if config_path: - with open(config_path, 'r', encoding='utf-8') as f: - config = json.load(f) - templates = config.get('templates', []) - else: - templates = [] - return templates - - except Exception as e: - logger.log_message(f"Error getting all available templates: {str(e)}", level=logging.ERROR) - return [] - -def toggle_user_template_preference(user_id, template_id, is_enabled, db_session): - """ - Toggle a user's template preference (enable/disable). - - Args: - user_id: ID of the user - template_id: ID of the template - is_enabled: Whether to enable or disable the template - db_session: Database session - - Returns: - Tuple (success: bool, message: str) - """ - try: - from src.db.schemas.models import UserTemplatePreference, AgentTemplate - from datetime import datetime, UTC - - # Verify template exists and is active - template = db_session.query(AgentTemplate).filter( - AgentTemplate.template_id == template_id, - AgentTemplate.is_active == True - ).first() - - if not template: - return False, "Template not found or inactive" - - # Check if preference record exists - preference = db_session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == template_id - ).first() - - if preference: - # Update existing preference - preference.is_enabled = is_enabled - preference.updated_at = datetime.now(UTC) - else: - # Create new preference record - preference = UserTemplatePreference( - user_id=user_id, - template_id=template_id, - is_enabled=is_enabled, - usage_count=0, - created_at=datetime.now(UTC), - updated_at=datetime.now(UTC) - ) - db_session.add(preference) - - db_session.commit() - - action = "enabled" if is_enabled else "disabled" - return True, f"Template '{template.template_name}' {action} successfully" - - except Exception as e: - db_session.rollback() - logger.log_message(f"Error toggling template preference: {str(e)}", level=logging.ERROR) - return False, f"Error updating template preference: {str(e)}" - - - -def load_all_available_templates_from_db(db_session): - """ - Load ALL available individual variant template agents from the database for direct access. - This allows users to use any individual template via @template_name regardless of preferences. - - Args: - db_session: Database session - - Returns: - Dict of template agent signatures keyed by template name - """ - try: - from src.db.schemas.models import AgentTemplate - - agent_signatures = {} - - # Get all active individual variant templates - all_templates = db_session.query(AgentTemplate).filter( - AgentTemplate.is_active == True, - AgentTemplate.variant_type.in_(['individual', 'both']) - ).all() - - for template in all_templates: - # Create dynamic signature for all active templates - signature = create_custom_agent_signature( - template.template_name, - template.description, - template.prompt_template, - template.category # Pass the category from database - ) - agent_signatures[template.template_name] = signature - - return agent_signatures - - except Exception as e: - logger.log_message(f"Error loading all available templates: {str(e)}", level=logging.ERROR) - return {} +logger = Logger("agents", see_time=True, console_log=False) +AGENTS_WITH_DESCRIPTION = { + "preprocessing_agent": "Cleans and prepares a DataFrame using Pandas and NumPy—handles missing values, detects column types, and converts date strings to datetime.", + "statistical_analytics_agent": "Performs statistical analysis (e.g., regression, seasonal decomposition) using statsmodels, with proper handling of categorical data and missing values.", + "sk_learn_agent": "Trains and evaluates machine learning models using scikit-learn, including classification, regression, and clustering with feature importance insights.", + "data_viz_agent": "Generates interactive visualizations with Plotly, selecting the best chart type to reveal trends, comparisons, and insights based on the analysis goal." +} -# === END CUSTOM AGENT FUNCTIONALITY === +PLANNER_AGENTS_WITH_DESCRIPTION = { + "planner_preprocessing_agent": ( + "Cleans and prepares a DataFrame using Pandas and NumPy—" + "handles missing values, detects column types, and converts date strings to datetime. " + "Outputs a cleaned DataFrame for the planner_statistical_analytics_agent." + ), + "planner_statistical_analytics_agent": ( + "Takes the cleaned DataFrame from preprocessing, performs statistical analysis " + "(e.g., regression, seasonal decomposition) using statsmodels with proper handling " + "of categorical data and remaining missing values. " + "Produces summary statistics and model diagnostics for the planner_sk_learn_agent." + ), + "planner_sk_learn_agent": ( + "Receives summary statistics and the cleaned data, trains and evaluates machine " + "learning models using scikit-learn (classification, regression, clustering), " + "and generates performance metrics and feature importance. " + "Passes the trained models and evaluation results to the planner_data_viz_agent." + ), + "planner_data_viz_agent": ( + "Consumes trained models and evaluation results to create interactive visualizations " + "with Plotly—selects the best chart type, applies styling, and annotates insights. " + "Delivers ready-to-share figures that communicate model performance and key findings." + ), +} def get_agent_description(agent_name, is_planner=False): - """ - Get agent description from database instead of hardcoded dictionaries. - This function is kept for backward compatibility but will fetch from DB. - """ - try: - from src.db.init_db import session_factory - from src.db.schemas.models import AgentTemplate - - db_session = session_factory() - try: - template = db_session.query(AgentTemplate).filter( - AgentTemplate.template_name == agent_name, - AgentTemplate.is_active == True - ).first() - - if template: - return template.description - else: - return "No description available for this agent" - finally: - db_session.close() - except Exception as e: - return "No description available for this agent" + if is_planner: + return PLANNER_AGENTS_WITH_DESCRIPTION[agent_name.lower()] if agent_name.lower() in PLANNER_AGENTS_WITH_DESCRIPTION else "No description available for this agent" + else: + return AGENTS_WITH_DESCRIPTION[agent_name.lower()] if agent_name.lower() in AGENTS_WITH_DESCRIPTION else "No description available for this agent" # Agent to make a Chat history name from a query @@ -390,165 +57,73 @@ class chat_history_name_agent(dspy.Signature): name = dspy.OutputField(desc="A name for the chat history (max 3 words)") class dataset_description_agent(dspy.Signature): - """ - - Generate a structured dataset context/description like this, you will be given headers for the data & existing description! -{ - "exact": "apple_stock_historical_data", - "description": "Daily historical stock market data for Apple Inc. including open, close, high, low prices, trading volume, and adjusted close for accurate return calculations.", - "columns": { - "Date": { - "type": "datetime", - "format": "YYYY-MM-DD", - "description": "Trading date", - "preprocessing": "Convert strings using pd.to_datetime(df['Date'], format='%Y-%m-%d')", - "missing_values_handling": "Interpolate or forward-fill missing dates for continuity" - }, - "Open": { - "type": "float", - "description": "Opening stock price in USD", - "preprocessing": "Direct float conversion" - }, - "High": { - "type": "float", - "description": "Highest stock price in USD during the trading day", - "preprocessing": "Direct float conversion" - }, - "Low": { - "type": "float", - "description": "Lowest stock price in USD during the trading day", - "preprocessing": "Direct float conversion" - }, - "Close": { - "type": "float", - "description": "Closing stock price in USD", - "preprocessing": "Direct float conversion" - }, - "Adj Close": { - "type": "float", - "description": "Adjusted closing price accounting for dividends and splits", - "preprocessing": "Direct float conversion" - }, - "Volume": { - "type": "integer", - "description": "Number of shares traded during the day", - "preprocessing": "Direct integer conversion" - } - }, - - - "usage_notes": "Ensure adjusted close prices are used for return calculations. Use consistent timezone if merging with other datasets. Handle missing values carefully to maintain temporal continuity.", - + """You are an AI agent that generates a detailed description of a given dataset for both users and analysis agents. +Your description should serve two key purposes: +1. Provide users with context about the dataset's purpose, structure, and key attributes. +2. Give analysis agents critical data handling instructions to prevent common errors. + +For data handling instructions, you must always include Python data types and address the following: +- Data type warnings (e.g., numeric columns stored as strings that need conversion). +- Null value handling recommendations. +- Format inconsistencies that require preprocessing. +- Explicit warnings about columns that appear numeric but are stored as strings (e.g., '10' vs 10). +- Explicit Python data types for each major column (e.g., int, float, str, bool, datetime). +- Columns with numeric values that should be treated as categorical (e.g., zip codes, IDs). +- Any date parsing or standardization required (e.g., MM/DD/YYYY to datetime). +- Any other technical considerations that would affect downstream analysis or modeling. +- List all columns and their data types with exact case sensitive spelling + +If an existing description is provided, enhance it with both business context and technical guidance for analysis agents, preserving accurate information from the existing description or what the user has written. + +Ensure the description is comprehensive and provides actionable insights for both users and analysis agents. + + +Example: +This housing dataset contains property details including price, square footage, bedrooms, and location data. +It provides insights into real estate market trends across different neighborhoods and property types. + +TECHNICAL CONSIDERATIONS FOR ANALYSIS: +- price (str): Appears numeric but is stored as strings with a '$' prefix and commas (e.g., "$350,000"). Requires cleaning with str.replace('$','').replace(',','') and conversion to float. +- square_footage (str): Contains unit suffix like 'sq ft' (e.g., "1,200 sq ft"). Remove suffix and commas before converting to int. +- bedrooms (int): Correctly typed but may contain null values (~5% missing) – consider imputation or filtering. +- zip_code (int): Numeric column but should be treated as str or category to preserve leading zeros and prevent unintended numerical analysis. +- year_built (float): May contain missing values (~15%) – consider mean/median imputation or exclusion depending on use case. +- listing_date (str): Dates stored in "MM/DD/YYYY" format – convert to datetime using pd.to_datetime(). +- property_type (str): Categorical column with inconsistent capitalization (e.g., "Condo", "condo", "CONDO") – normalize to lowercase for consistent grouping. """ dataset = dspy.InputField(desc="The dataset to describe, including headers, sample data, null counts, and data types.") existing_description = dspy.InputField(desc="An existing description to improve upon (if provided).", default="") - description = dspy.OutputField(desc="A comprehensive dataset context with business context and technical guidance for analysis agents.") + description = dspy.OutputField(desc="A comprehensive dataset description with business context and technical guidance for analysis agents.") -class custom_agent_instruction_generator(dspy.Signature): - """You are an AI agent instruction generator that creates comprehensive, professional prompts for custom data analysis agents. - - Your task is to take a user's requirements and generate a detailed agent instruction that follows the same structure and quality as the default system agents (preprocessing_agent, statistical_analytics_agent, sk_learn_agent, data_viz_agent). - - Key requirements for generated instructions: - 1. **Professional Structure**: Use clear sections with headers and bullet points - 2. **Input/Output Specification**: Clearly define what the agent receives and produces - 3. **Technical Guidelines**: Include specific library recommendations and best practices - 4. **Error Handling**: Include instructions for handling common issues - 5. **Code Quality**: Emphasize clean, reproducible, and well-documented code - 6. **Standardized Outputs**: Ensure consistent format with 'code' and 'summary' outputs - - Structure your instructions as follows: - - Brief role definition and purpose - - Input specifications and expectations - - Core responsibilities and tasks - - Technical requirements and best practices - - Library and methodology recommendations - - Error handling and edge cases - - Output format requirements - - Example code patterns (if relevant) - - Categories and their focus areas: - - **Visualization**: - - Emphasize Plotly for interactive charts - - Include styling and layout best practices - - Focus on chart type selection based on data - - Performance optimization for large datasets - - Color schemes and accessibility - - **Modelling**: - - Cover model selection and evaluation - - Include cross-validation and metrics - - Emphasize feature engineering and preprocessing - - Handle different problem types (classification, regression, clustering) - - Include hyperparameter tuning guidance - - **Data Manipulation**: - - Focus on Pandas and NumPy operations - - Include data cleaning and transformation - - Handle missing values and outliers - - Emphasize data type conversions - - Include aggregation and reshaping operations - - Make instructions generic enough to handle various tasks within the category while being specific enough to provide clear guidance. - Always include the standard output format: 'code' (Python code) and 'summary' (bullet-point explanation). - - Example instruction structure: - ''' - You are a [specific role] agent specializing in [category focus]. Your task is to [main purpose]... - - **Input Requirements:** - - dataset: [description] - - goal: [description] - - [other inputs as needed] - - **Core Responsibilities:** - 1. [Primary task] - 2. [Secondary task] - 3. [Additional requirements] - - **Technical Guidelines:** - - Use [recommended libraries] - - Follow [best practices] - - Handle [common issues] - - **Output Requirements:** - - code: [code specification] - - summary: [summary specification] - ''' - """ - category = dspy.InputField(desc="The category of the custom agent: 'Visualization', 'Modelling', or 'Data Manipulation'") - user_requirements = dspy.InputField(desc="User's description of what they want the agent to do, including specific tasks, methods, or focus areas") - agent_instructions = dspy.OutputField(desc="Complete, professional agent instructions following the structure and quality of default system agents, ready to be used as a custom agent prompt") - -class advanced_query_planner(dspy.Signature): +class analytical_planner(dspy.Signature): """ -You are a advanced data analytics planner agent. Your task is to generate the most efficient plan—using the fewest necessary agents and variables—to achieve a user-defined goal. The plan must preserve data integrity, avoid unnecessary steps, and ensure clear data flow between agents. - -**CRITICAL**: Before planning, check if any agents are available in Agent_desc. If Agent_desc is empty or contains no active agents, respond with: -plan: no_agents_available -plan_instructions: {"message": "No agents are currently enabled for analysis. Please enable at least one agent (preprocessing, statistical analysis, machine learning, or visualization) in your template preferences to proceed with data analysis."} +You are a data analytics planner agent. Your task is to generate the most efficient plan—using the fewest necessary agents and variables—to achieve a user-defined goal. The plan must preserve data integrity, avoid unnecessary steps, and ensure clear data flow between agents. **Inputs**: + 1. Datasets (raw or preprocessed) 2. Agent descriptions (roles, variables they create/use, constraints) 3. User-defined goal (e.g., prediction, analysis, visualization) + **Responsibilities**: + 1. **Feasibility**: Confirm the goal is achievable with the provided data and agents; ask for clarification if it's unclear. 2. **Minimal Plan**: Use the smallest set of agents and variables; avoid redundant transformations; ensure agents are ordered logically and only included if essential. 3. **Instructions**: For each agent, define: + * **create**: output variables and their purpose * **use**: input variables and their role - * **instruction**: concise explanation of the agent's function and relevance to the goal + * **instruction**: concise explanation of the agent’s function and relevance to the goal 4. **Clarity**: Keep instructions precise; avoid intermediate steps unless necessary; ensure each agent has a distinct, relevant role. + ### **Output Format**: Example: 1 agent use goal: "Generate a bar plot showing sales by category after cleaning the raw data and calculating the average of the 'sales' column" Output: - plan: data_viz_agent + plan: planner_data_viz_agent { - "data_viz_agent": { + "planner_data_viz_agent": { "create": [ "cleaned_data: DataFrame - cleaned version of df (pd.Dataframe) after removing null values" ], @@ -558,11 +133,13 @@ Output: "instruction": "Clean df by removing null values, calculate the average sales, and generate a bar plot showing sales by category." } } + Example 3 Agent goal:"Clean the dataset, run a linear regression to model the relationship between marketing budget and sales, and visualize the regression line with confidence intervals." -plan: preprocessing_agent -> statistical_analytics_agent -> data_viz_agent +plan: planner_preprocessing_agent -> planner_statistical_analytics_agent -> planner_data_viz_agent + { - "preprocessing_agent": { + "planner_preprocessing_agent": { "create": [ "cleaned_data: DataFrame - cleaned version of df with missing values handled and proper data types inferred" ], @@ -571,7 +148,8 @@ plan: preprocessing_agent -> statistical_analytics_agent -> data_viz_agent ], "instruction": "Clean df by handling missing values and converting column types (e.g., dates). Output cleaned_data for modeling." }, - "statistical_analytics_agent": { + + "planner_statistical_analytics_agent": { "create": [ "regression_results: dict - model summary including coefficients, p-values, R², and confidence intervals" ], @@ -580,7 +158,8 @@ plan: preprocessing_agent -> statistical_analytics_agent -> data_viz_agent ], "instruction": "Perform linear regression using cleaned_data to model sales as a function of marketing budget. Return regression_results including coefficients and confidence intervals." }, - "data_viz_agent": { + + "planner_data_viz_agent": { "create": [ "regression_plot: PlotlyFigure - visual plot showing regression line with confidence intervals" ], @@ -591,249 +170,52 @@ plan: preprocessing_agent -> statistical_analytics_agent -> data_viz_agent "instruction": "Generate a Plotly regression plot using cleaned_data and regression_results. Show the fitted line, scatter points, and 95% confidence intervals." } } -Try to use as few agents to answer the user query as possible. -Respond in the user's language for all explanations and instructions, but keep all code, variable names, function names, model names, agent names, and library names in English. - """ - dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df, columns set df as copy of df") - Agent_desc = dspy.InputField(desc="The agents available in the system") - goal = dspy.InputField(desc="The user defined goal") - plan = dspy.OutputField(desc="The plan that would achieve the user defined goal", prefix='Plan:') - plan_instructions = dspy.OutputField(desc="Detailed variable-level instructions per agent for the plan") - -class basic_query_planner(dspy.Signature): - """ - You are the basic query planner in the system, you pick one agent, to answer the user's goal. - Use the Agent_desc that describes the names and actions of agents available. - - **CRITICAL**: Before planning, check if any agents are available in Agent_desc. If Agent_desc is empty or contains no active agents, respond with: - plan: no_agents_available - plan_instructions: {"message": "No agents are currently enabled for analysis. Please enable at least one agent (preprocessing, statistical analysis, machine learning, or visualization) in your template preferences to proceed with data analysis."} - - Example: Visualize height and salary? - plan:data_viz_agent - plan_instructions: - { - "data_viz_agent": { - "create": ["scatter_plot"], - "use": ["original_data"], - "instruction": "use the original_data to create scatter_plot of height & salary, using plotly" - } - } - Example: Tell me the correlation between X and Y - plan:preprocessing_agent - plan_instructions:{ - "data_viz_agent": { - "create": ["correlation"], - "use": "use": ["original_data"], - "instruction": "use the original_data to measure correlation of X & Y, using pandas" - } - - - Respond in the user's language for all explanations and instructions, but keep all code, variable names, function names, model names, agent names, and library names in English. - original_data is placeholder, use exact_python_name: name_of_df for actual dataset name +Keep it as simple as possible, unless user specifies indepth query! """ dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df, columns set df as copy of df") Agent_desc = dspy.InputField(desc="The agents available in the system") goal = dspy.InputField(desc="The user defined goal") - plan = dspy.OutputField(desc="The plan that would achieve the user defined goal", prefix='Plan:') - plan_instructions = dspy.OutputField(desc="Instructions on what the agent should do alone") - - -class intermediate_query_planner(dspy.Signature): - # The planner agent which routes the query to Agent(s) - # The output is like this Agent1->Agent2 etc - """ You are an intermediate data analytics planner agent. You have access to three inputs - 1. Datasets - 2. Data Agent descriptions - 3. User-defined Goal - You take these three inputs to develop a comprehensive plan to achieve the user-defined goal from the data & Agents available. - In case you think the user-defined goal is infeasible you can ask the user to redefine or add more description to the goal. - - **CRITICAL**: Before planning, check if any agents are available in Agent_desc. If Agent_desc is empty or contains no active agents, respond with: - plan: no_agents_available - plan_instructions: {"message": "No agents are currently enabled for analysis. Please enable at least one agent (preprocessing, statistical analysis, machine learning, or visualization) in your template preferences to proceed with data analysis."} - - Give your output in this format: - plan: Agent1->Agent2 - plan_instructions = { - "Agent1": { - "create": ["aggregated_variable"], - "use": ["original_data"] - "instruction": "use the original_data to create aggregated_variable" - }, - "Agent2": { - "create": ["visualization_of_data"], - "use": ["aggregated_variable,original_data"], - "instruction": "use the aggregated_variable & original_data to create visualization_of_data" - } - } - Keep the instructions minimal without many variables, and minimize the number of unknowns, keep it obvious! - Try to use no more than 2 agents, unless completely necessary! - original_data is placeholder, use exact_python_name: name_of_df for actual dataset name - - - - Respond in the user's language for all explanations and instructions, but keep all code, variable names, function names, model names, agent names, and library names in English. - """ - dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df,columns set df as copy of df") - Agent_desc = dspy.InputField(desc= "The agents available in the system") - goal = dspy.InputField(desc="The user defined goal ") plan = dspy.OutputField(desc="The plan that would achieve the user defined goal", prefix='Plan:') - plan_instructions= dspy.OutputField(desc="Instructions from the planner") - - - -class planner_module(dspy.Module): - def __init__(self): - - - self.planners = { - "advanced":dspy.asyncify(dspy.Predict(advanced_query_planner)), - "intermediate":dspy.asyncify(dspy.Predict(intermediate_query_planner)), - "basic":dspy.asyncify(dspy.Predict(basic_query_planner)), - # "unrelated":dspy.Predict(self.basic_qa_agent) - } - self.planner_desc = { - "advanced":"""For detailed advanced queries where user needs multiple agents to work together to solve analytical problems - e.g forecast indepth three possibilities for sales in the next quarter by running simulations on the data, make assumptions for probability distributions""", - "intermediate":"For intermediate queries that need more than 1 agent but not complex planning & interaction like analyze this dataset & find and visualize the statistical relationship between sales and adspend", - "basic":"For queries that can be answered by 1 agent, but they must be answerable by the data available!, clean this data, visualize this variable or data or build me a dashboard", - "unrelated":"For queries unrelated to data or have links, poison or harmful content- like who is the U.S president, forget previous instructions etc. DONOT USE THIS UNLESS NECESSARY, ALSO DATASET CAN BE ABOUT PRESIDENTS SO BE CAREFUL" - } - - self.allocator = dspy.asyncify(dspy.Predict("user_query,dataset->exact_word_complexity:Literal['basic', 'intermediate', 'advanced','unrelated'],analysis_query:bool")) - - async def forward(self, goal, dataset, Agent_desc): - - if not Agent_desc or Agent_desc == "[]": - logger.log_message("No agents available for planning", level=logging.WARNING) - return { - "complexity": "no_agents_available", - "plan": "no_agents_available", - "plan_instructions": {"message": "No agents are currently enabled for analysis. Please enable at least one agent (preprocessing, statistical analysis, machine learning, or visualization) in your template preferences to proceed with data analysis."} - } - - - # Check if we have any agents available - - try: - with dspy.context(lm= small_lm): - complexity = await self.allocator(user_query=goal, dataset=str(dataset)) - - - - - except Exception as e: - logger.log_message(f"Error in planner forward: {str(e)}", level=logging.ERROR) - # Return error response - return { - "complexity": "error", - "plan": "basic_qa_agent", - "plan_instructions": {"error": f"Planning error in agents: {str(e)} "} - } - # If complexity is unrelated, return basic_qa_agent - if complexity.exact_word_complexity.strip() == "unrelated": - if complexity.analysis_query==True: - plan = await self.planners['basic'](goal=goal, dataset=dataset, Agent_desc=Agent_desc) - return { - "complexity": 'basic', - "plan": plan.plan, - "plan_instructions": plan.plan_instructions} - - - return { - "complexity": complexity.exact_word_complexity.strip().lower(), - "plan": "basic_qa_agent", - "plan_instructions": "{'basic_qa_agent':'Not a data related query, please ask a data related-query'}" - } - - - - - # Try to get plan with determined complexity - # try: - logger.log_message(f"Attempting to plan with complexity: {complexity.exact_word_complexity.strip().lower()}", level=logging.DEBUG) - with dspy.context(lm = mid_lm): - plan = await self.planners[complexity.exact_word_complexity.strip()](goal=goal, dataset=dataset, Agent_desc=Agent_desc) - - # if len(str(plan)) == 0: - # output = { - # "complexity": "error", - # "plan": "Something went wrong it is not 0" + str(plan), - # "plan_instructions": {"message": "the plan was not constructed" + str(Agent_desc)} - # } - # else: - # output = { - # "complexity": "error", - # "plan": "Something went wrong it is 0" + str(plan), - # "plan_instructions": {"message": "the plan was not constructed" + str(Agent_desc)} - # } - - - - - # If plan or plan.plan is not available, return an error response - if not plan or not hasattr(plan, 'plan'): - logger.log_message("Planner did not return a valid plan object or 'plan' attribute is missing", level=logging.ERROR) - return { - "complexity": complexity.exact_word_complexity.strip().lower(), - "plan": "planning_error", - "plan_instructions": {"error": "Planner did not return a valid plan. Please try again or check agent configuration."} - } - - logger.log_message(f"Plan generated successfully: {plan}", level=logging.DEBUG) - - # Check if the planner returned no_agents_available - if hasattr(plan, 'plan') and 'no_agents_available' in str(plan.plan): - logger.log_message("Planner returned no_agents_available", level=logging.WARNING) - output = { - "complexity": "no_agents_available", - "plan": "no_agents_available", - "plan_instructions": {"message": "No agents are currently enabled for analysis. Please enable at least one agent (preprocessing, statistical analysis, machine learning, or visualization) in your template preferences to proceed with data analysis."} - } + plan_instructions = dspy.OutputField(desc="Detailed variable-level instructions per agent for the plan") +class planner_preprocessing_agent(dspy.Signature): + """ +You are a preprocessing agent in a multi-agent data analytics system. - output = { - "complexity": complexity.exact_word_complexity.strip().lower(), - "plan": plan.plan, - "plan_instructions": plan.plan_instructions - } - - return output +You are given: +* A dataset (already loaded as `df`). +* A user-defined analysis goal (e.g., predictive modeling, exploration, cleaning). +* Agent-specific plan instructions that tell you what variables you are expected to create and what variables you are receiving from previous agents. +* processed_df is just an arbitrary name, it can be anything the planner says to clean! +### Your Responsibilities: +* Follow the provided plan and create only the required variables listed in the 'create' section of the plan instructions. +* Do not create fake data or introduce variables not explicitly part of the instructions. +* Do not read data from CSV ; the dataset (`df`) is already loaded and ready for processing. +* Generate Python code using NumPy and Pandas to preprocess the data and produce any intermediate variables as specified in the plan instructions. +### Best Practices for Preprocessing: -class preprocessing_agent(dspy.Signature): - """ -You are a preprocessing agent that can work both individually and in multi-agent data analytics systems. -You are given: -* A dataset (already loaded as with exact_python_name mentioned). -* A user-defined analysis goal (e.g., predictive modeling, exploration, cleaning). -* Optional plan instructions that tell you what variables you are expected to create and what variables you are receiving from previous agents. - -### Your Responsibilities: -* If plan_instructions are provided, follow the provided plan and create only the required variables listed in the 'create' section. -* If no plan_instructions are provided, perform standard data preprocessing based on the goal. -* Do not create fake data or introduce variables not explicitly part of the instructions. -* Do not read data from CSV; the dataset (`df`) is already loaded and ready for processing. -* Generate Python code using NumPy and Pandas to preprocess the data and produce any intermediate variables as specified. +1. Create a copy of the original DataFrame : It will always be stored as df, it already exists use it! -### Best Practices for Preprocessing: -1. Create a copy of the original DataFrame: It will always be stored as df, it already exists use it! ```python processed_df = df.copy() ``` -2. Separate column types: + +2. Separate column types : + ```python numeric_cols = processed_df.select_dtypes(include='number').columns categorical_cols = processed_df.select_dtypes(include='object').columns ``` -3. Handle missing values: + +3. Handle missing values : + ```python for col in numeric_cols: processed_df[col] = processed_df[col].fillna(processed_df[col].median()) @@ -841,7 +223,9 @@ You are given: for col in categorical_cols: processed_df[col] = processed_df[col].fillna(processed_df[col].mode()[0] if not processed_df[col].mode().empty else 'Unknown') ``` -4. Convert string columns to datetime safely: + +4. Convert string columns to datetime safely : + ```python def safe_to_datetime(x): try: @@ -851,144 +235,200 @@ You are given: cleaned_df['date_column'] = cleaned_df['date_column'].apply(safe_to_datetime) ``` -5. Do not alter the DataFrame index unless explicitly instructed. -6. Log assumptions and corrections in comments to clarify any choices made during preprocessing. -7. Do not mutate global state: Avoid in-place modifications unless clearly necessary (e.g., using `.copy()`). -8. Handle data types properly: + +> Replace `processed_df`,'cleaned_df' and `date_column` with whatever names the user or planner provides. + + +5. Do not alter the DataFrame index : + Avoid using `reset_index()`, `set_index()`, or reindexing unless explicitly instructed. + +6. Log assumptions and corrections in comments to clarify any choices made during preprocessing. + +7. Do not mutate global state : Avoid in-place modifications unless clearly necessary (e.g., using `.copy()`). + +8. Handle data types properly : + * Avoid coercing types blindly (e.g., don't compare timestamps to strings or floats). * Use `pd.to_datetime(..., errors='coerce')` for safe datetime parsing. -9. Preserve column structure: Only drop or rename columns if explicitly instructed. + +9. Preserve column structure : Only drop or rename columns if explicitly instructed. ### Output: -1. Code: Python code that performs the requested preprocessing steps. -2. Summary: A brief explanation of what preprocessing was done (e.g., columns handled, missing value treatment). + +1. Code : Python code that performs the requested preprocessing steps as per the plan instructions. +2. Summary : A brief explanation of what preprocessing was done (e.g., columns handled, missing value treatment). ### Principles to Follow: -- Never alter the DataFrame index unless explicitly instructed. -- Handle missing data explicitly, filling with default values when necessary. -- Preserve column structure and avoid unnecessary modifications. -- Ensure data types are appropriate (e.g., dates parsed correctly). -- Log assumptions in the code. -Respond in the user's language for all summary and reasoning but keep the code in english + +-Never alter the DataFrame index unless explicitly instructed. +-Handle missing data explicitly, filling with default values when necessary. +-Preserve column structure and avoid unnecessary modifications. +-Ensure data types are appropriate (e.g., dates parsed correctly). +-Log assumptions in the code. + """ dataset = dspy.InputField(desc="The dataset, preloaded as df") goal = dspy.InputField(desc="User-defined goal for the analysis") - plan_instructions = dspy.InputField(desc="Agent-level instructions about what to create and receive (optional for individual use)", default="") + plan_instructions = dspy.InputField(desc="Agent-level instructions about what to create and receive") code = dspy.OutputField(desc="Generated Python code for preprocessing") summary = dspy.OutputField(desc="Explanation of what was done and why") - -class data_viz_agent(dspy.Signature): +class planner_data_viz_agent(dspy.Signature): """ -You are a data visualization agent that can work both individually and in multi-agent analytics pipelines. -Your primary responsibility is to generate visualizations based on the user-defined goal. + ### **Data Visualization Agent Definition** + + You are the **data visualization agent** in a multi-agent analytics pipeline. Your primary responsibility is to **generate visualizations** based on the **user-defined goal** and the **plan instructions**. You are provided with: + * **goal**: A user-defined goal outlining the type of visualization the user wants (e.g., "plot sales over time with trendline"). -* **dataset**: The dataset (e.g., `df_cleaned`) which will be passed to you by other agents in the pipeline. Do not assume or create any variables — the data is already present and valid when you receive it. + * **dataset**: The dataset (e.g., `df_cleaned`) which will be passed to you by other agents in the pipeline. **Do not assume or create any variables** — **the data is already present and valid** when you receive it. * **styling_index**: Specific styling instructions (e.g., axis formatting, color schemes) for the visualization. -* **plan_instructions**: Optional dictionary containing: - * **'create'**: List of visualization components you must generate (e.g., 'scatter_plot', 'bar_chart'). - * **'use'**: List of variables you must use to generate the visualizations. - * **'instructions'**: Additional instructions related to the creation of the visualizations. + * **plan_instructions**: A dictionary containing: + + * **'create'**: List of **visualization components** you must generate (e.g., 'scatter_plot', 'bar_chart'). + * **'use'**: List of **variables you must use** to generate the visualizations. This includes datasets and any other variables provided by the other agents. + * **'instructions'**: A list of additional instructions related to the creation of the visualizations, such as requests for trendlines or axis formats. + + --- + + ### **Responsibilities**: -### Responsibilities: 1. **Strict Use of Provided Variables**: - * You must never create fake data. Only use the variables and datasets that are explicitly provided. - * If plan_instructions are provided and any variable listed in plan_instructions['use'] is missing, return an error. - * If no plan_instructions are provided, work with the available dataset directly. + + * You must **never create fake data**. Only use the variables and datasets that are explicitly **provided** to you in the `plan_instructions['use']` section. All the required data **must already be available**. + * If any variable listed in `plan_instructions['use']` is missing or invalid, **you must return an error** and not proceed with any visualization. 2. **Visualization Creation**: - * Based on the goal and optional 'create' section of plan_instructions, generate the required visualization using Plotly. - * Respect the user-defined goal in determining which type of visualization to create. + + * Based on the **'create'** section of the `plan_instructions`, generate the **required visualization** using **Plotly**. For example, if the goal is to plot a time series, you might generate a line chart. + * Respect the **user-defined goal** in determining which type of visualization to create. 3. **Performance Optimization**: - * If the dataset contains more than 50,000 rows, you must sample the data to 5,000 rows to improve performance: + + * If the dataset contains **more than 50,000 rows**, you **must sample** the data to **5,000 rows** to improve performance. Use this method: + ```python if len(df) > 50000: df = df.sample(5000, random_state=42) ``` 4. **Layout and Styling**: - * Apply formatting and layout adjustments as defined by the styling_index. - * Ensure that all axes (x and y) have consistent formats (e.g., using `K`, `M`, or 1,000 format, but not mixing formats). + + * Apply formatting and layout adjustments as defined by the **styling_index**. This may include: + + * Axis labels and title formatting. + * Tick formats for axes. + * Color schemes or color maps for visual elements. + * You must ensure that all axes (x and y) have **consistent formats** (e.g., using `K`, `M`, or 1,000 format, but not mixing formats). 5. **Trendlines**: - * Trendlines should only be included if explicitly requested in the goal or plan_instructions. + + * Trendlines should **only be included** if explicitly requested in the **'instructions'** section of `plan_instructions`. 6. **Displaying the Visualization**: + * Use Plotly's `fig.show()` method to display the created chart. - * Never output raw datasets or the goal itself. Only the visualization code and the chart should be returned. + * **Never** output raw datasets or the **goal** itself. Only the visualization code and the chart should be returned. 7. **Error Handling**: - * If required dataset or variables are missing, return an error message indicating which specific variable is missing. - * If the goal or create instructions are ambiguous, return an error stating the issue. + + * If the required dataset or variables are missing or invalid (i.e., not included in `plan_instructions['use']`), return an error message indicating which specific variable is missing or invalid. + * If the **goal** or **create** instructions are ambiguous or invalid, return an error stating the issue. 8. **No Data Modification**: - * Never modify the provided dataset or generate new data. If the data needs preprocessing, assume it's already been done by other agents. - -### Important Notes: -- Use update_yaxes, update_xaxes, not axis -- Each visualization must be generated as a separate figure using go.Figure() -- Do NOT use subplots under any circumstances -- Each figure must be returned individually using: fig.to_html(full_html=False) -- Use update_layout with xaxis and yaxis only once per figure -- Enhance readability with low opacity (0.4-0.7) where appropriate -- Apply visually distinct colors for different elements or categories -- Use only one number format consistently: either 'K', 'M', or comma-separated values -- Only include trendlines in scatter plots if the user explicitly asks for them -- Always end each visualization with: fig.to_html(full_html=False) - -Respond in the user's language for all summary and reasoning but keep the code in english + + * **Never** modify the provided dataset or generate new data. If the data needs preprocessing or cleaning, assume it's already been done by other agents. + + --- + + ### **Strict Conditions**: + + * You **never** create any data. + * You **only** use the data and variables passed to you. + * If any required data or variable is missing or invalid, **you must stop** and return a clear error message. + + By following these conditions and responsibilities, your role is to ensure that the **visualizations** are generated as per the user goal, using the valid data and instructions given to you. + """ goal = dspy.InputField(desc="User-defined chart goal (e.g. trendlines, scatter plots)") dataset = dspy.InputField(desc="Details of the dataframe (`df`) and its columns") styling_index = dspy.InputField(desc="Instructions for plot styling and layout formatting") - plan_instructions = dspy.InputField(desc="Variables to create and receive for visualization purposes (optional for individual use)", default="") + plan_instructions = dspy.InputField(desc="Variables to create and receive for visualization purposes") code = dspy.OutputField(desc="Plotly Python code for the visualization") summary = dspy.OutputField(desc="Plain-language summary of what is being visualized") -class statistical_analytics_agent(dspy.Signature): +class planner_statistical_analytics_agent(dspy.Signature): """ -You are a statistical analytics agent that can work both individually and in multi-agent data analytics pipelines. +**Agent Definition:** + +You are a statistical analytics agent in a multi-agent data analytics pipeline. + You are given: + * A dataset (usually a cleaned or transformed version like `df_cleaned`). * A user-defined goal (e.g., regression, seasonal decomposition). -* Optional plan instructions specifying: - * Which variables you are expected to CREATE (e.g., `regression_model`). - * Which variables you will USE (e.g., `df_cleaned`, `target_variable`). - * A set of instructions outlining additional processing or handling for these variables. +* Agent-specific **plan instructions** specifying: + + * Which **variables** you are expected to **CREATE** (e.g., `regression_model`). + * Which **variables** you will **USE** (e.g., `df_cleaned`, `target_variable`). + * A set of **instructions** outlining additional processing or handling for these variables (e.g., handling missing values, adding constants, transforming features, etc.). + +**Your Responsibilities:** -### Your Responsibilities: * Use the `statsmodels` library to implement the required statistical analysis. * Ensure that all strings are handled as categorical variables via `C(col)` in model formulas. * Always add a constant using `sm.add_constant()`. -* Do not modify the DataFrame's index. +* Do **not** modify the DataFrame's index. * Convert `X` and `y` to float before fitting the model. * Handle missing values before modeling. * Avoid any data visualization (that is handled by another agent). * Write output to the console using `print()`. -### If the goal is regression: +**If the goal is regression:** + * Use `statsmodels.OLS` with proper handling of categorical variables and adding a constant term. * Handle missing values appropriately. -### If the goal is seasonal decomposition: +**If the goal is seasonal decomposition:** + * Use `statsmodels.tsa.seasonal_decompose`. * Ensure the time series and period are correctly provided (i.e., `period` should not be `None`). -### Instructions to Follow: -1. If plan_instructions are provided: - * CREATE only the variables specified in plan_instructions['CREATE']. Do not create any intermediate or new variables. - * USE only the variables specified in plan_instructions['USE'] to carry out the task. - * Follow any additional instructions in plan_instructions['INSTRUCTIONS']. - * Do not reassign or modify any variables passed via plan_instructions. -2. If no plan_instructions are provided, perform standard statistical analysis based on the goal and available data. +**You must not:** + +* You must always create the variables in `plan_instructions['CREATE']`. +* **Never create the `df` variable**. Only work with the variables passed via the `plan_instructions`. +* Rely on hardcoded column names — use those passed via `plan_instructions`. +* Introduce or modify intermediate variables unless they are explicitly listed in `plan_instructions['CREATE']`. + +**Instructions to Follow:** + +1. **CREATE** only the variables specified in `plan_instructions['CREATE']`. Do not create any intermediate or new variables. +2. **USE** only the variables specified in `plan_instructions['USE']` to carry out the task. +3. Follow any **additional instructions** in `plan_instructions['INSTRUCTIONS']` (e.g., preprocessing steps, encoding, handling missing values). +4. **Do not reassign or modify** any variables passed via `plan_instructions`. These should be used as-is. + +**Example Workflow:** +Given that the `plan_instructions` specifies variables to **CREATE** and **USE**, and includes instructions, your approach should look like this: + +1. Use `df_cleaned` and the variables like `X` and `y` from `plan_instructions` for analysis. +2. Follow instructions for preprocessing (e.g., handle missing values or scale features). +3. If the goal is regression: + + * Use `sm.OLS` for model fitting. + * Handle categorical variables via `C(col)` and add a constant term. +4. If the goal is seasonal decomposition: + + * Ensure `period` is provided and use `sm.tsa.seasonal_decompose`. +5. Store the output variable as specified in `plan_instructions['CREATE']`. ### Example Code Structure: + ```python import statsmodels.api as sm + def statistical_model(X, y, goal, period=None): try: X = X.dropna() @@ -1000,104 +440,162 @@ def statistical_model(X, y, goal, period=None): # Add constant term to X X = sm.add_constant(X) + if goal == 'regression': formula = 'y ~ ' + ' + '.join([f'C({col})' if X[col].dtype.name == 'category' else col for col in X.columns]) model = sm.OLS(y.astype(float), X.astype(float)).fit() - return model.summary() + regression_model = model.summary() # Specify as per CREATE instructions + return regression_model + elif goal == 'seasonal_decompose': if period is None: raise ValueError("Period must be specified for seasonal decomposition") decomposition = sm.tsa.seasonal_decompose(y, period=period) - return decomposition + seasonal_decomposition = decomposition # Specify as per CREATE instructions + return seasonal_decomposition + else: - raise ValueError("Unknown goal specified. Please provide a valid goal.") + raise ValueError("Unknown goal specified.") except Exception as e: return f"An error occurred: {e}" ``` -### Summary: -1. Always USE the variables passed in plan_instructions['USE'] to carry out the task (if provided). -2. Only CREATE the variables specified in plan_instructions['CREATE'] (if provided). -3. Follow any additional instructions in plan_instructions['INSTRUCTIONS'] (if provided). +**Summary:** + +1. Always **USE** the variables passed in `plan_instructions['USE']` to carry out the task. +2. Only **CREATE** the variables specified in `plan_instructions['CREATE']`. Do not create any additional variables. +3. Follow any **additional instructions** in `plan_instructions['INSTRUCTIONS']` (e.g., handling missing values, adding constants). 4. Ensure reproducibility by setting the random state appropriately and handling categorical variables. 5. Focus on statistical analysis and avoid any unnecessary data manipulation. -### Output: -* The code implementing the statistical analysis, including all required steps. -* A summary of what the statistical analysis does, how it's performed, and why it fits the goal. -* Respond in the user's language for all summary and reasoning but keep the code in english +**Output:** + +* The **code** implementing the statistical analysis, including all required steps. +* A **summary** of what the statistical analysis does, how it's performed, and why it fits the goal. + """ dataset = dspy.InputField(desc="Preprocessed dataset, often named df_cleaned") goal = dspy.InputField(desc="The user's statistical analysis goal, e.g., regression or seasonal_decompose") - plan_instructions = dspy.InputField(desc="Instructions on variables to create and receive for statistical modeling (optional for individual use)", default="") + plan_instructions = dspy.InputField(desc="Instructions on variables to create and receive for statistical modeling") code = dspy.OutputField(desc="Python code for statistical modeling using statsmodels") - summary = dspy.OutputField(desc="A concise bullet-point summary of the statistical analysis performed and key findings") + summary = dspy.OutputField(desc="Explanation of statistical analysis steps") -class sk_learn_agent(dspy.Signature): + +class planner_sk_learn_agent(dspy.Signature): """ -You are a machine learning agent that can work both individually and in multi-agent data analytics pipelines. + **Agent Definition:** + + You are a machine learning agent in a multi-agent data analytics pipeline. + You are given: + * A dataset (often cleaned and feature-engineered). * A user-defined goal (e.g., classification, regression, clustering). -* Optional plan instructions specifying: - * Which variables you are expected to CREATE (e.g., `trained_model`, `predictions`). - * Which variables you will USE (e.g., `df_cleaned`, `target_variable`, `feature_columns`). - * A set of instructions outlining additional processing or handling for these variables. + * Agent-specific **plan instructions** specifying: + + * Which **variables** you are expected to **CREATE** (e.g., `trained_model`, `predictions`). + * Which **variables** you will **USE** (e.g., `df_cleaned`, `target_variable`, `feature_columns`). + * A set of **instructions** outlining additional processing or handling for these variables (e.g., handling missing values, applying transformations, or other task-specific guidelines). + + **Your Responsibilities:** -### Your Responsibilities: * Use the scikit-learn library to implement the appropriate ML pipeline. * Always split data into training and testing sets where applicable. * Use `print()` for all outputs. * Ensure your code is: - * Reproducible: Set `random_state=42` wherever applicable. - * Modular: Avoid deeply nested code. - * Focused on model building, not visualization (leave plotting to the `data_viz_agent`). + + * **Reproducible**: Set `random_state=42` wherever applicable. + * **Modular**: Avoid deeply nested code. + * **Focused on model building**, not visualization (leave plotting to the `data_viz_agent`). * Your task may include: + * Preprocessing inputs (e.g., encoding). * Model selection and training. * Evaluation (e.g., accuracy, RMSE, classification report). -### You must not: - * Visualize anything (that's another agent's job). -* Rely on hardcoded column names — use those passed via plan_instructions or infer from data. -* Never create or modify any variables not explicitly mentioned in plan_instructions['CREATE'] (if provided). -* Never create the `df` variable. You will only work with the variables passed via the plan_instructions. -* Do not introduce intermediate variables unless they are listed in plan_instructions['CREATE'] (if provided). - -### Instructions to Follow: -1. If plan_instructions are provided: - * CREATE only the variables specified in the plan_instructions['CREATE'] list. - * USE only the variables specified in the plan_instructions['USE'] list. - * Follow any processing instructions in the plan_instructions['INSTRUCTIONS'] list. - * Do not reassign or modify any variables passed via plan_instructions. -2. If no plan_instructions are provided, perform standard machine learning analysis based on the goal and available data. - -### Example Workflow: -Given that the plan_instructions specifies variables to CREATE and USE, and includes instructions, your approach should look like this: -1. Use `df_cleaned` and `feature_columns` from the plan_instructions to extract your features (`X`). -2. Use `target_column` from plan_instructions to extract your target (`y`). - 3. If instructions are provided (e.g., scale or encode), follow them. - 4. Split data into training and testing sets using `train_test_split`. - 5. Train the model based on the received goal (classification, regression, etc.). -6. Store the output variables as specified in plan_instructions['CREATE']. + **You must not:** + + * Visualize anything (that's another agent's job). + * Rely on hardcoded column names — use those passed via `plan_instructions`. + * **Never create or modify any variables not explicitly mentioned in `plan_instructions['CREATE']`.** + * **Never create the `df` variable**. You will **only** work with the variables passed via the `plan_instructions`. + * Do not introduce intermediate variables unless they are listed in `plan_instructions['CREATE']`. + + **Instructions to Follow:** + + 1. **CREATE** only the variables specified in the `plan_instructions['CREATE']` list. Do not create any intermediate or new variables. + 2. **USE** only the variables specified in the `plan_instructions['USE']` list. You are **not allowed** to create or modify any variables not listed in the plan instructions. + 3. Follow any **processing instructions** in the `plan_instructions['INSTRUCTIONS']` list. This might include tasks like handling missing values, scaling features, or encoding categorical variables. Always perform these steps on the variables specified in the `plan_instructions`. + 4. Do **not reassign or modify** any variables passed via `plan_instructions`. These should be used as-is. + + **Example Workflow:** + Given that the `plan_instructions` specifies variables to **CREATE** and **USE**, and includes instructions, your approach should look like this: + + 1. Use `df_cleaned` and `feature_columns` from the `plan_instructions` to extract your features (`X`). + 2. Use `target_column` from `plan_instructions` to extract your target (`y`). + 3. If instructions are provided (e.g., scale or encode), follow them. + 4. Split data into training and testing sets using `train_test_split`. + 5. Train the model based on the received goal (classification, regression, etc.). + 6. Store the output variables as specified in `plan_instructions['CREATE']`. + + ### Example Code Structure: + + ```python + from sklearn.model_selection import train_test_split + from sklearn.linear_model import LogisticRegression + from sklearn.metrics import classification_report + from sklearn.preprocessing import StandardScaler + + # Ensure that all variables follow plan instructions: + # Use received inputs: df_cleaned, feature_columns, target_column + X = df_cleaned[feature_columns] + y = df_cleaned[target_column] + + # Apply any preprocessing instructions (e.g., scaling if instructed) + if 'scale' in plan_instructions['INSTRUCTIONS']: + scaler = StandardScaler() + X = scaler.fit_transform(X) + + # Split the data into training and testing sets + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + + # Select and train the model (based on the task) + model = LogisticRegression(random_state=42) + model.fit(X_train, y_train) + + # Generate predictions + predictions = model.predict(X_test) + + # Create the variable specified in 'plan_instructions': 'metrics' + metrics = classification_report(y_test, predictions) + + # Print the results + print(metrics) + + # Ensure the 'metrics' variable is returned as requested in the plan + ``` + + **Summary:** -### Summary: -1. Always USE the variables passed in plan_instructions['USE'] to build the pipeline (if provided). -2. Only CREATE the variables specified in plan_instructions['CREATE'] (if provided). -3. Follow any additional instructions in plan_instructions['INSTRUCTIONS'] (if provided). -4. Ensure reproducibility by setting random_state=42 wherever necessary. + 1. Always **USE** the variables passed in `plan_instructions['USE']` to build the pipeline. + 2. Only **CREATE** the variables specified in `plan_instructions['CREATE']`. Do not create any additional variables. + 3. Follow any **additional instructions** in `plan_instructions['INSTRUCTIONS']` (e.g., preprocessing steps). + 4. Ensure reproducibility by setting `random_state=42` wherever necessary. 5. Focus on model building, evaluation, and saving the required outputs—avoid any unnecessary variables. -### Output: -* The code implementing the ML task, including all required steps. -* A summary of what the model does, how it is evaluated, and why it fits the goal. - * Respond in the user's language for all summary and reasoning but keep the code in english + **Output:** + + * The **code** implementing the ML task, including all required steps. + * A **summary** of what the model does, how it is evaluated, and why it fits the goal. + + + """ dataset = dspy.InputField(desc="Input dataset, often cleaned and feature-selected (e.g., df_cleaned)") goal = dspy.InputField(desc="The user's machine learning goal (e.g., classification or regression)") - plan_instructions = dspy.InputField(desc="Instructions indicating what to create and what variables to receive (optional for individual use)", default="") + plan_instructions = dspy.InputField(desc="Instructions indicating what to create and what variables to receive") code = dspy.OutputField(desc="Scikit-learn based machine learning code") summary = dspy.OutputField(desc="Explanation of the ML approach and evaluation") @@ -1111,7 +609,167 @@ class goal_refiner_agent(dspy.Signature): goal = dspy.InputField(desc="The user defined goal ") refined_goal = dspy.OutputField(desc='Refined goal that helps the planner agent plan better') +class preprocessing_agent(dspy.Signature): + """You are a AI data-preprocessing agent. Generate clean and efficient Python code using NumPy and Pandas to perform introductory data preprocessing on a pre-loaded DataFrame df, based on the user's analysis goals. + + Preprocessing Requirements: + + 1. Identify Column Types + - Separate columns into numeric and categorical using: + categorical_columns = df.select_dtypes(include=[object, 'category']).columns.tolist() + numeric_columns = df.select_dtypes(include=[np.number]).columns.tolist() + + 2. Handle Missing Values + - Numeric columns: Impute missing values using the mean of each column + - Categorical columns: Impute missing values using the mode of each column + + 3. Convert Date Strings to Datetime + - For any column suspected to represent dates (in string format), convert it to datetime using: + def safe_to_datetime(date): + try: + return pd.to_datetime(date, errors='coerce', cache=False) + except (ValueError, TypeError): + return pd.NaT + df['datetime_column'] = df['datetime_column'].apply(safe_to_datetime) + - Replace 'datetime_column' with the actual column names containing date-like strings + + Important Notes: + - Do NOT create a correlation matrix — correlation analysis is outside the scope of preprocessing + - Do NOT generate any plots or visualizations + + Output Instructions: + 1. Include the full preprocessing Python code + 2. Provide a brief bullet-point summary of the steps performed. Example: + • Identified 5 numeric and 4 categorical columns + • Filled missing numeric values with column means + • Filled missing categorical values with column modes + • Converted 1 date column to datetime format + """ + dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df, column_names set df as copy of df") + goal = dspy.InputField(desc="The user defined goal could ") + code = dspy.OutputField(desc ="The code that does the data preprocessing and introductory analysis") + summary = dspy.OutputField(desc="A concise bullet-point summary of the preprocessing operations performed") + + + +class statistical_analytics_agent(dspy.Signature): + # Statistical Analysis Agent, builds statistical models using StatsModel Package + """ + You are a statistical analytics agent. Your task is to take a dataset and a user-defined goal and output Python code that performs the appropriate statistical analysis to achieve that goal. Follow these guidelines: + + IMPORTANT: You may be provided with previous interaction history. The section marked "### Current Query:" contains the user's current request. Any text in "### Previous Interaction History:" is for context only and is NOT part of the current request. + + Data Handling: + + Always handle strings as categorical variables in a regression using statsmodels C(string_column). + Do not change the index of the DataFrame. + Convert X and y into float when fitting a model. + Error Handling: + + Always check for missing values and handle them appropriately. + Ensure that categorical variables are correctly processed. + Provide clear error messages if the model fitting fails. + Regression: + + For regression, use statsmodels and ensure that a constant term is added to the predictor using sm.add_constant(X). + Handle categorical variables using C(column_name) in the model formula. + Fit the model with model = sm.OLS(y.astype(float), X.astype(float)).fit(). + Seasonal Decomposition: + + Ensure the period is set correctly when performing seasonal decomposition. + Verify the number of observations works for the decomposition. + Output: + + Ensure the code is executable and as intended. + Also choose the correct type of model for the problem + Avoid adding data visualization code. + + Use code like this to prevent failing: + import pandas as pd + import numpy as np + import statsmodels.api as sm + + def statistical_model(X, y, goal, period=None): + try: + # Check for missing values and handle them + X = X.dropna() + y = y.loc[X.index].dropna() + + # Ensure X and y are aligned + X = X.loc[y.index] + # Convert categorical variables + for col in X.select_dtypes(include=['object', 'category']).columns: + X[col] = X[col].astype('category') + + # Add a constant term to the predictor + X = sm.add_constant(X) + + # Fit the model + if goal == 'regression': + # Handle categorical variables in the model formula + formula = 'y ~ ' + ' + '.join([f'C({col})' if X[col].dtype.name == 'category' else col for col in X.columns]) + model = sm.OLS(y.astype(float), X.astype(float)).fit() + return model.summary() + + elif goal == 'seasonal_decompose': + if period is None: + raise ValueError("Period must be specified for seasonal decomposition") + decomposition = sm.tsa.seasonal_decompose(y, period=period) + return decomposition + + else: + raise ValueError("Unknown goal specified. Please provide a valid goal.") + + except Exception as e: + return f"An error occurred: {e}" + + # Example usage: + result = statistical_analysis(X, y, goal='regression') + print(result) + + If visualizing use plotly + + Provide a concise bullet-point summary of the statistical analysis performed. + + Example Summary: + • Applied linear regression with OLS to predict house prices based on 5 features + • Model achieved R-squared of 0.78 + • Significant predictors include square footage (p<0.001) and number of bathrooms (p<0.01) + • Detected strong seasonal pattern with 12-month periodicity + • Forecast shows 15% growth trend over next quarter + + """ + dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df,columns set df as copy of df") + goal = dspy.InputField(desc="The user defined goal for the analysis to be performed") + code = dspy.OutputField(desc ="The code that does the statistical analysis using statsmodel") + summary = dspy.OutputField(desc="A concise bullet-point summary of the statistical analysis performed and key findings") + + +class sk_learn_agent(dspy.Signature): + # Machine Learning Agent, performs task using sci-kit learn + """You are a machine learning agent. + Your task is to take a dataset and a user-defined goal, and output Python code that performs the appropriate machine learning analysis to achieve that goal. + You should use the scikit-learn library. + + IMPORTANT: You may be provided with previous interaction history. The section marked "### Current Query:" contains the user's current request. Any text in "### Previous Interaction History:" is for context only and is NOT part of the current request. + + Make sure your output is as intended! + + Provide a concise bullet-point summary of the machine learning operations performed. + + Example Summary: + • Trained a Random Forest classifier on customer churn data with 80/20 train-test split + • Model achieved 92% accuracy and 88% F1-score + • Feature importance analysis revealed that contract length and monthly charges are the strongest predictors of churn + • Implemented K-means clustering (k=4) on customer shopping behaviors + • Identified distinct segments: high-value frequent shoppers (22%), occasional big spenders (35%), budget-conscious regulars (28%), and rare visitors (15%) + + """ + dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df,columns. set df as copy of df") + goal = dspy.InputField(desc="The user defined goal ") + code = dspy.OutputField(desc ="The code that does the Exploratory data analysis") + summary = dspy.OutputField(desc="A concise bullet-point summary of the machine learning analysis performed and key results") @@ -1126,12 +784,16 @@ class code_combiner_agent(dspy.Signature): # Combines code from different agents into one script """ You are a code combine agent, taking Python code output from many agents and combining the operations into 1 output You also fix any errors in the code. + IMPORTANT: You may be provided with previous interaction history. The section marked "### Current Query:" contains the user's current request. Any text in "### Previous Interaction History:" is for context only and is NOT part of the current request. + Double check column_names/dtypes using dataset, also check if applied logic works for the datatype df = df.copy() Also add this to display Plotly chart fig.show() + Make sure your output is as intended! + Provide a concise bullet-point summary of the code integration performed. Example Summary: @@ -1139,7 +801,7 @@ class code_combiner_agent(dspy.Signature): • Fixed variable scope issues, standardized DataFrame handling (e.g., using `df.copy()`), and corrected errors. • Validated column names and data types against the dataset definition to prevent runtime issues. • Ensured visualizations are displayed correctly (e.g., added `fig.show()`). - Respond in the user's language for all summary and reasoning but keep the code in english + """ dataset = dspy.InputField(desc="Use this double check column_names, data types") agent_code_list =dspy.InputField(desc="A list of code given by each agent") @@ -1147,16 +809,78 @@ class code_combiner_agent(dspy.Signature): summary = dspy.OutputField(desc="A concise 4 bullet-point summary of the code integration performed and improvements made") +class data_viz_agent(dspy.Signature): + # Visualizes data using Plotly + """ + You are an AI agent responsible for generating interactive data visualizations using Plotly. + + IMPORTANT Instructions: + + - The section marked "### Current Query:" contains the user's request. Any text in "### Previous Interaction History:" is for context only and should NOT be treated as part of the current request. + - You must only use the tools provided to you. This agent handles visualization only. + - If len(df) > 50000, always sample the dataset before visualization using: + if len(df) > 50000: + df = df.sample(50000, random_state=1) + + - Each visualization must be generated as a **separate figure** using go.Figure(). + Do NOT use subplots under any circumstances. + + - Each figure must be returned individually using: + fig.to_html(full_html=False) + + - Use update_layout with xaxis and yaxis **only once per figure**. + + - Enhance readability and clarity by: + • Using low opacity (0.4-0.7) where appropriate + • Applying visually distinct colors for different elements or categories + + - Make sure the visual **answers the user's specific goal**: + • Identify what insight or comparison the user is trying to achieve + • Choose the visualization type and features (e.g., color, size, grouping) to emphasize that goal + • For example, if the user asks for "trends in revenue," use a time series line chart; if they ask for "top-performing categories," use a bar chart sorted by value + • Prioritize highlighting patterns, outliers, or comparisons relevant to the question + + - Never include the dataset or styling index in the output. + + - If there are no relevant columns for the requested visualization, respond with: + "No relevant columns found to generate this visualization." + + - Use only one number format consistently: either 'K', 'M', or comma-separated values like 1,000/1,000,000. Do not mix formats. + + - Only include trendlines in scatter plots if the user explicitly asks for them. + + - Output only the code and a concise bullet-point summary of what the visualization reveals. + + - Always end each visualization with: + fig.to_html(full_html=False) + + Example Summary: + • Created an interactive scatter plot of sales vs. marketing spend with color-coded product categories + • Included a trend line showing positive correlation (r=0.72) + • Highlighted outliers where high marketing spend resulted in low sales + • Generated a time series chart of monthly revenue from 2020-2023 + • Added annotations for key business events + • Visualization reveals 35% YoY growth with seasonal peaks in Q4 + """ + goal = dspy.InputField(desc="user defined goal which includes information about data and chart they want to plot") + dataset = dspy.InputField(desc=" Provides information about the data in the data frame. Only use column names and dataframe_name as in this context") + styling_index = dspy.InputField(desc='Provides instructions on how to style your Plotly plots') + code= dspy.OutputField(desc="Plotly code that visualizes what the user needs according to the query & dataframe_index & styling_context") + summary = dspy.OutputField(desc="A concise bullet-point summary of the visualization created and key insights revealed") + + class code_fix(dspy.Signature): """ You are an expert AI developer and data analyst assistant, skilled at identifying and resolving issues in Python code related to data analytics. Another agent has attempted to generate Python code for a data analytics task but produced code that is broken or throws an error. + Your task is to: 1. Carefully examine the provided **faulty_code** and the corresponding **error** message. 2. Identify the **exact cause** of the failure based on the error and surrounding context. 3. Modify only the necessary portion(s) of the code to fix the issue, utilizing the **dataset_context** to inform your corrections. 4. Ensure the **intended behavior** of the original code is preserved (e.g., if the code is meant to filter, group, or visualize data, that functionality must be preserved). 5. Ensure the final output is **runnable**, **error-free**, and **logically consistent**. + Strict instructions: - Assume the dataset is already loaded and available in the code context; do not include any code to read, load, or create data. - Do **not** modify any working parts of the code unnecessarily. @@ -1164,27 +888,32 @@ Strict instructions: - Do **not** output anything besides the corrected, full version of the code (i.e., no explanations, comments, or logs). - Avoid introducing new dependencies or libraries unless absolutely required to fix the problem. - The output must be complete and executable as-is. + Be precise, minimal, and reliable. Prioritize functional correctness. + One-shot example: === dataset_context: "This dataset contains historical price and trading data for two major financial assets: the S&P 500 index and Bitcoin (BTC). The data includes daily price metrics (open, high, low, close) and percentage changes for both assets... Change % columns are stored as strings with '%' symbol (e.g., '-5.97%') and require cleaning." + faulty_code: # Convert percentage strings to floats df['Change %'] = df['Change %'].str.rstrip('%').astype(float) df['Change % BTC'] = df['Change % BTC'].str.rstrip('%').astype(float) + error: Error in data_viz_agent: Can only use .str accessor with string values! Traceback (most recent call last): File "/app/scripts/format_response.py", line 196, in execute_code_from_markdown exec(block_code, context) AttributeError: Can only use .str accessor with string values! + fixed_code: # Convert percentage strings to floats df['Change %'] = df['Change %'].astype(str).str.rstrip('%').astype(float) df['Change % BTC'] = df['Change % BTC'].astype(str).str.rstrip('%').astype(float) -Respond in the user's language for all summary and reasoning but keep the code in english === + """ dataset_context = dspy.InputField(desc="The dataset context to be used for the code fix") faulty_code = dspy.InputField(desc="The faulty Python code used for data analytics") @@ -1194,17 +923,20 @@ Respond in the user's language for all summary and reasoning but keep the code i class code_edit(dspy.Signature): """ You are an expert AI code editor that specializes in modifying existing data analytics code based on user requests. The user provides a working or partially working code snippet, a natural language prompt describing the desired change, and dataset context information. + Your job is to: 1. Analyze the provided original_code, user_prompt, and dataset_context. 2. Modify only the part(s) of the code that are relevant to the user's request, using the dataset context to inform your edits. 3. Leave all unrelated parts of the code unchanged, unless the user explicitly requests a full rewrite or broader changes. 4. Ensure that your changes maintain or improve the functionality and correctness of the code. + Strict requirements: - Assume the dataset is already loaded and available in the code context; do not include any code to read, load, or create data. - Do not change variable names, function structures, or logic outside the scope of the user's request. - Do not refactor, optimize, or rewrite unless explicitly instructed. - Ensure the edited code remains complete and executable. - Output only the modified code, without any additional explanation, comments, or extra formatting. + Make your edits precise, minimal, and faithful to the user's instructions, using the dataset context to guide your modifications. """ dataset_context = dspy.InputField(desc="The dataset context to be used for the code edit, including information about the dataset's shape, columns, types, and null values") @@ -1212,366 +944,101 @@ Make your edits precise, minimal, and faithful to the user's instructions, using user_prompt = dspy.InputField(desc="The user instruction describing how the code should be changed") edited_code = dspy.OutputField(desc="The updated version of the code reflecting the user's request, incorporating changes informed by the dataset context") - - - - - # The ind module is called when agent_name is # explicitly mentioned in the query class auto_analyst_ind(dspy.Module): """Handles individual agent execution when explicitly specified in query""" - def __init__(self, agents, retrievers, user_id=None, db_session=None): + def __init__(self, agents, retrievers): # Initialize agent modules and retrievers self.agents = {} self.agent_inputs = {} self.agent_desc = [] - # logger.log_message(f"[INIT] Initializing auto_analyst_ind with user_id={user_id}, agents={len(agents) if agents else 0}", level=logging.INFO) - - # Load core agents based on user preferences (not always loaded) - if not agents and user_id and db_session: - try: - # Get user preferences for core agents - from src.db.schemas.models import AgentTemplate, UserTemplatePreference - - core_agent_names = ['preprocessing_agent', 'statistical_analytics_agent', 'sk_learn_agent', 'data_viz_agent'] - - for agent_name in core_agent_names: - logger.log_message(f"[INIT] Processing core agent: {agent_name}", level=logging.DEBUG) - - # Check if user has enabled this core agent - template = db_session.query(AgentTemplate).filter( - AgentTemplate.template_name == agent_name, - AgentTemplate.is_active == True - ).first() - - if not template: - logger.log_message(f"[INIT] Core agent template '{agent_name}' not found in database", level=logging.WARNING) - continue - - # Get the agent signature class - if agent_name == 'preprocessing_agent': - agent_signature = preprocessing_agent - elif agent_name == 'statistical_analytics_agent': - agent_signature = statistical_analytics_agent - elif agent_name == 'sk_learn_agent': - agent_signature = sk_learn_agent - elif agent_name == 'data_viz_agent': - agent_signature = data_viz_agent - - # Add to agents dict - self.agents[agent_name] = dspy.asyncify(dspy.ChainOfThought(agent_signature)) - - # Set input fields based on signature - if agent_name == 'data_viz_agent': - self.agent_inputs[agent_name] = {'goal', 'dataset', 'styling_index', 'plan_instructions'} - else: - self.agent_inputs[agent_name] = {'goal', 'dataset', 'plan_instructions'} - - # Get description from database - self.agent_desc.append({agent_name: get_agent_description(agent_name)}) - # logger.log_message(f"[INIT] Successfully loaded core agent: {agent_name} with inputs: {self.agent_inputs[agent_name]}", level=logging.INFO) - - except Exception as e: - logger.log_message(f"[INIT] Error loading core agents based on preferences: {str(e)}", level=logging.ERROR) - # Fallback to loading all core agents if preference system fails - self._load_default_agents_fallback() - elif not agents: - self._load_default_agents_fallback() - # If no user_id/db_session provided, load all core agents as fallback - # logger.log_message(f"[INIT] No agents provided and no user_id/db_session, loading fallback agents", level=logging.INFO) + # Create modules from agent signatures + for i, a in enumerate(agents): + name = a.__pydantic_core_schema__['schema']['model_name'] + self.agents[name] = dspy.ChainOfThoughtWithHint(a) + self.agent_inputs[name] = {x.strip() for x in str(agents[i].__pydantic_core_schema__['cls']).split('->')[0].split('(')[1].split(',')} + self.agent_desc.append(get_agent_description(name)) - - # Load ALL available template agents if user_id and db_session are provided - # For individual agent execution (@agent_name), users should be able to access any available agent - if user_id and db_session: - try: - # For individual use, load ALL available templates regardless of user preferences - template_signatures = load_all_available_templates_from_db(db_session) - - # logger.log_message(f"[INIT] Loaded {len(template_signatures)} template signatures from database", level=logging.INFO) - - for template_name, signature in template_signatures.items(): - # Skip if this is a core agent - we'll load it separately - if template_name in ['preprocessing_agent', 'statistical_analytics_agent', 'sk_learn_agent', 'data_viz_agent']: - # logger.log_message(f"[INIT] Skipping template {template_name} as it's a core agent", level=logging.DEBUG) - continue - - # Add template agent to agents dict - self.agents[template_name] = dspy.asyncify(dspy.ChainOfThought(signature)) - - # Determine if this is a visualization agent based on database category - is_viz_agent = False - try: - from src.db.schemas.models import AgentTemplate - - # Find template record to check category - template_record = db_session.query(AgentTemplate).filter( - AgentTemplate.template_name == template_name - ).first() - - if template_record and template_record.category and template_record.category.lower() == 'visualization': - is_viz_agent = True - else: - # Fallback to name-based detection for legacy templates - is_viz_agent = ('viz' in template_name.lower() or - 'visual' in template_name.lower() or - 'plot' in template_name.lower() or - 'chart' in template_name.lower() or - 'matplotlib' in template_name.lower()) - except Exception as cat_error: - logger.log_message(f"[INIT] Error checking category for template {template_name}: {str(cat_error)}", level=logging.WARNING) - # Fallback to name-based detection - is_viz_agent = ('viz' in template_name.lower() or - 'visual' in template_name.lower() or - 'plot' in template_name.lower() or - 'chart' in template_name.lower() or - 'matplotlib' in template_name.lower()) - - # Set input fields based on agent type - if is_viz_agent: - self.agent_inputs[template_name] = {'goal', 'dataset', 'styling_index', 'plan_instructions'} - else: - self.agent_inputs[template_name] = {'goal', 'dataset', 'plan_instructions'} - - # Store template agent description - try: - if not template_record: - template_record = db_session.query(AgentTemplate).filter( - AgentTemplate.template_name == template_name - ).first() - - if template_record: - description = f"Template: {template_record.description}" - self.agent_desc.append({template_name: description}) - else: - self.agent_desc.append({template_name: f"Template: {template_name}"}) - except Exception as desc_error: - logger.log_message(f"[INIT] Error getting description for template {template_name}: {str(desc_error)}", level=logging.WARNING) - self.agent_desc.append({template_name: f"Template: {template_name}"}) - - # logger.log_message(f"[INIT] Successfully loaded template agent: {template_name} with inputs: {self.agent_inputs[template_name]}, is_viz_agent: {is_viz_agent}", level=logging.INFO) - - except Exception as e: - logger.log_message(f"[INIT] Error loading template agents for user {user_id}: {str(e)}", level=logging.ERROR) - - self.agents['basic_qa_agent'] = dspy.asyncify(dspy.Predict("goal->answer")) - self.agent_inputs['basic_qa_agent'] = {"goal"} - self.agent_desc.append({'basic_qa_agent':"Answers queries unrelated to data & also that include links, poison or attempts to attack the system"}) - - # Initialize retrievers (no planner needed for individual agent execution) - self.dataset = retrievers['dataframe_index'] + # Initialize components + self.memory_summarize_agent = dspy.ChainOfThought(m.memory_summarize_agent) + self.dataset = retrievers['dataframe_index'].as_retriever(k=1) self.styling_index = retrievers['style_index'].as_retriever(similarity_top_k=1) + self.code_combiner_agent = dspy.ChainOfThought(code_combiner_agent) - # Store user_id for usage tracking - self.user_id = user_id - - # Log final summary - # logger.log_message(f"[INIT] Initialization complete. Total agents loaded: {len(self.agents)}", level=logging.INFO) - # logger.log_message(f"[INIT] Available agents: {list(self.agents.keys())}", level=logging.INFO) - # logger.log_message(f"[INIT] Agent inputs mapping: {self.agent_inputs}", level=logging.DEBUG) - - def _load_default_agents_fallback(self): - """Fallback method to load default agents when preference system fails""" - # logger.log_message("Loading default agents as fallback for auto_analyst_ind", level=logging.WARNING) - - # Load the 4 core agents from database - core_agent_names = ['preprocessing_agent', 'statistical_analytics_agent', 'sk_learn_agent', 'data_viz_agent'] - - for agent_name in core_agent_names: - # Get the agent signature class - if agent_name == 'preprocessing_agent': - agent_signature = preprocessing_agent - elif agent_name == 'statistical_analytics_agent': - agent_signature = statistical_analytics_agent - elif agent_name == 'sk_learn_agent': - agent_signature = sk_learn_agent - elif agent_name == 'data_viz_agent': - agent_signature = data_viz_agent - - # Add to agents dict - self.agents[agent_name] = dspy.asyncify(dspy.ChainOfThought(agent_signature)) - - # Set input fields based on signature - if agent_name == 'data_viz_agent': - self.agent_inputs[agent_name] = {'goal', 'dataset', 'styling_index', 'plan_instructions'} - else: - self.agent_inputs[agent_name] = {'goal', 'dataset', 'plan_instructions'} - - # Get description from database - self.agent_desc.append({agent_name: get_agent_description(agent_name)}) - # logger.log_message(f"Added fallback agent: {agent_name}", level=logging.DEBUG) - - async def _track_agent_usage(self, agent_name): - """Track usage for template agents""" - try: - # Skip tracking for standard agents - if agent_name in ['preprocessing_agent', 'statistical_analytics_agent', 'sk_learn_agent', 'data_viz_agent', 'basic_qa_agent']: - return - - # Only track if we have user_id (template agents) - if not self.user_id: - return - - from src.db.init_db import session_factory - from src.db.schemas.models import AgentTemplate, UserTemplatePreference - from datetime import datetime, UTC - - # Create database session - session = session_factory() - try: - # Find the template - template = session.query(AgentTemplate).filter( - AgentTemplate.template_name == agent_name - ).first() - - if not template: - logger.log_message(f"Template '{agent_name}' not found for usage tracking", level=logging.WARNING) - return - - # Find or create user template preference record - preference = session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == self.user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - if not preference: - # Create new preference record (disabled by default) - preference = UserTemplatePreference( - user_id=self.user_id, - template_id=template.template_id, - is_enabled=False, # Disabled by default - usage_count=0, - last_used_at=None, - created_at=datetime.now(UTC), - updated_at=datetime.now(UTC) - ) - session.add(preference) - - # Update usage tracking - preference.usage_count += 1 - preference.last_used_at = datetime.now(UTC) - preference.updated_at = datetime.now(UTC) - session.commit() - - logger.log_message( - f"Tracked usage for template '{agent_name}' (count: {preference.usage_count})", - level=logging.DEBUG - ) - - except Exception as e: - session.rollback() - logger.log_message(f"Error tracking usage for template {agent_name}: {str(e)}", level=logging.ERROR) - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error in _track_agent_usage for {agent_name}: {str(e)}", level=logging.ERROR) + # Initialize thread pool + self.executor = ThreadPoolExecutor(max_workers=min(4, os.cpu_count() * 2)) - async def execute_agent(self, specified_agent, inputs): + def execute_agent(self, specified_agent, inputs): """Execute agent and generate memory summary in parallel""" try: - # logger.log_message(f"[EXECUTE] Starting execution of agent: {specified_agent}", level=logging.INFO) - # logger.log_message(f"[EXECUTE] Agent inputs: {inputs}", level=logging.DEBUG) - # Execute main agent - agent_result = await self.agents[specified_agent.strip()](**inputs) - - # Track usage for custom agents and templates - await self._track_agent_usage(specified_agent.strip()) - - # logger.log_message(f"[EXECUTE] Agent {specified_agent} execution completed successfully", level=logging.INFO) + agent_result = self.agents[specified_agent.strip()](**inputs) return specified_agent.strip(), dict(agent_result) except Exception as e: - # logger.log_message(f"[EXECUTE] Error executing agent {specified_agent}: {str(e)}", level=logging.ERROR) - - # logger.log_message(f"[EXECUTE] Full traceback: {traceback.format_exc()}", level=logging.ERROR) return specified_agent.strip(), {"error": str(e)} - async def forward(self, query, specified_agent): + def execute_agent_with_memory(self, specified_agent, inputs, query): + """Execute agent and generate memory summary in parallel""" try: - # logger.log_message(f"[FORWARD] Processing query with specified agent: {specified_agent}", level=logging.INFO) - # logger.log_message(f"[FORWARD] Query: {query}", level=logging.DEBUG) + # Execute main agent + agent_result = self.agents[specified_agent.strip()](**inputs) + agent_dict = dict(agent_result) + # Generate memory summary + memory_result = self.memory_summarize_agent( + agent_response=specified_agent+' '+agent_dict['code']+'\n'+agent_dict['summary'], + user_goal=query + ) + + return { + specified_agent.strip(): agent_dict, + 'memory_'+specified_agent.strip(): str(memory_result.summary) + } + except Exception as e: + return {"error": str(e)} + + def forward(self, query, specified_agent): + try: # If specified_agent contains multiple agents separated by commas # This is for handling multiple @agent mentions in one query if "," in specified_agent: agent_list = [agent.strip() for agent in specified_agent.split(",")] - # logger.log_message(f"[FORWARD] Multiple agents detected: {agent_list}", level=logging.INFO) - return await self.execute_multiple_agents(query, agent_list) + return self.execute_multiple_agents(query, agent_list) # Process query with specified agent (single agent case) dict_ = {} - dict_['dataset'] = self.dataset + dict_['dataset'] = self.dataset.retrieve(query)[0].text dict_['styling_index'] = self.styling_index.retrieve(query)[0].text - dict_['hint'] = [] dict_['goal'] = query dict_['Agent_desc'] = str(self.agent_desc) - - if specified_agent.strip() not in self.agent_inputs: - return {"response": f"Agent '{specified_agent.strip()}' not found in agent inputs"} - - # Create inputs that match exactly what the agent expects - inputs = {} - required_fields = self.agent_inputs[specified_agent.strip()] - - for field in required_fields: - if field == 'goal': - inputs['goal'] = query - elif field == 'dataset': - inputs['dataset'] = dict_['dataset'] - elif field == 'styling_index': - inputs['styling_index'] = dict_['styling_index'] - elif field == 'plan_instructions': - inputs['plan_instructions'] = "" # Empty for individual agent use - elif field == 'hint': - inputs['hint'] = "" # Empty string for hint - else: - # For any other fields, try to get from dict_ if available - if field in dict_: - inputs[field] = dict_[field] - else: - inputs[field] = "" # Provide empty string as fallback - - - if specified_agent.strip() not in self.agents: - return {"response": f"Agent '{specified_agent.strip()}' not found in agents"} - - result = await self.agents[specified_agent.strip()](**inputs) - - # Track usage for template agents - await self._track_agent_usage(specified_agent.strip()) - - try: - result_dict = dict(result) - except Exception as dict_error: - return {"response": f"Error converting agent result to dict: {str(dict_error)}"} + + # Prepare inputs + inputs = {x:dict_[x] for x in self.agent_inputs[specified_agent.strip()]} + inputs['hint'] = str(dict_['hint']).replace('[','').replace(']','') - output_dict = {specified_agent.strip(): result_dict} + # Execute agent + result = self.agents[specified_agent.strip()](**inputs) + output_dict = {specified_agent.strip(): dict(result)} - # Check for errors in the agent's response (not in the outer dict) - if "error" in result_dict: - return {"response": f"Error executing agent: {result_dict['error']}"} + if "error" in output_dict: + return {"response": f"Error executing agent: {output_dict['error']}"} return output_dict except Exception as e: - import traceback - logger.log_message(f"[FORWARD] Full traceback: {traceback.format_exc()}", level=logging.ERROR) return {"response": f"This is the error from the system: {str(e)}"} - async def execute_multiple_agents(self, query, agent_list): + def execute_multiple_agents(self, query, agent_list): """Execute multiple agents sequentially on the same query""" try: - logger.log_message(f"[MULTI] Executing multiple agents: {agent_list}", level=logging.INFO) - # Initialize resources dict_ = {} - dict_['dataset'] = self.dataset + dict_['dataset'] = self.dataset.retrieve(query)[0].text dict_['styling_index'] = self.styling_index.retrieve(query)[0].text dict_['hint'] = [] dict_['goal'] = query @@ -1582,559 +1049,152 @@ class auto_analyst_ind(dspy.Module): # Execute each agent sequentially for agent_name in agent_list: - logger.log_message(f"[MULTI] Processing agent: {agent_name}", level=logging.INFO) - if agent_name not in self.agents: - logger.log_message(f"[MULTI] Agent '{agent_name}' not found", level=logging.ERROR) results[agent_name] = {"error": f"Agent '{agent_name}' not found"} continue - # Create inputs that match exactly what the agent expects - inputs = {} - required_fields = self.agent_inputs[agent_name] - - logger.log_message(f"[MULTI] Required fields for {agent_name}: {required_fields}", level=logging.DEBUG) - - for field in required_fields: - if field == 'goal': - inputs['goal'] = query - elif field == 'dataset': - inputs['dataset'] = dict_['dataset'] - elif field == 'styling_index': - inputs['styling_index'] = dict_['styling_index'] - elif field == 'plan_instructions': - inputs['plan_instructions'] = "" # Empty for individual agent use - elif field == 'hint': - inputs['hint'] = "" # Empty string for hint - else: - # For any other fields, try to get from dict_ if available - if field in dict_: - inputs[field] = dict_[field] - else: - # logger.log_message(f"[MULTI] WARNING: Field '{field}' required by agent but not available in dict_", level=logging.WARNING) - pass - - # logger.log_message(f"[MULTI] Prepared inputs for {agent_name}: {list(inputs.keys())}", level=logging.DEBUG) + # Prepare inputs for this agent + inputs = {x:dict_[x] for x in self.agent_inputs[agent_name] if x in dict_} + inputs['hint'] = str(dict_['hint']).replace('[','').replace(']','') # Execute agent - try: - agent_result = await self.agents[agent_name](**inputs) - agent_dict = dict(agent_result) - results[agent_name] = agent_dict - - # Track usage for template agents - await self._track_agent_usage(agent_name) - - # Collect code for later combination - if 'code' in agent_dict: - code_list.append(agent_dict['code']) - - # logger.log_message(f"[MULTI] Successfully executed agent: {agent_name}", level=logging.INFO) - - except Exception as agent_error: - # logger.log_message(f"[MULTI] Error executing agent {agent_name}: {str(agent_error)}", level=logging.ERROR) - results[agent_name] = {"error": str(agent_error)} + agent_result = self.agents[agent_name](**inputs) + agent_dict = dict(agent_result) + results[agent_name] = agent_dict + + # Collect code for later combination + if 'code' in agent_dict: + code_list.append(agent_dict['code']) - # logger.log_message(f"[MULTI] Completed multiple agent execution. Results: {list(results.keys())}", level=logging.INFO) return results except Exception as e: - logger.log_message(f"[MULTI] Error executing multiple agents: {str(e)}", level=logging.ERROR) return {"response": f"Error executing multiple agents: {str(e)}"} -class data_context_gen(dspy.Signature): - """ - Generate a Python-friendly JSON structure that describes one or more datasets - loaded from Excel or CSV files. This helps the system understand the dataset - structure, semantics, and use cases. - - The JSON should include: - - Dataset name and source (file or sheet) - - Dataset role (transactional or reference) - - Description or business purpose - - Column names with: - - Data type (string, int, float, date, etc.) - - Semantic role: identifier, attribute, category, measure, temporal - - Relationships to other datasets (optional, natural-language style) - - Common metrics (as formulas or derived fields) - - Example use cases - - Example format: - { - "datasets": { - "sales_data": { - "source": "Sales_Data.csv", - "role": "transactional", - "description": "Sales transactions across regions and products.", - "columns": { - "order_id": {"type": "string", "role": "identifier"}, - "order_date": {"type": "date", "role": "temporal"}, - "region": {"type": "string", "role": "category"}, - "product_id": {"type": "string", "role": "identifier"}, - "quantity": {"type": "int", "role": "measure"}, - "unit_price": {"type": "float", "role": "measure"} - }, - "metrics": [ - "revenue = quantity * unit_price" - ], - "use_cases": [ - "Revenue trend analysis", - "Regional sales comparison" - ] - } - } - } - - Column roles: identifier, attribute, category, measure, temporal - Dataset roles: transactional, reference - - """ - user_description = dspy.InputField(desc="User's description of the data, including relationships") - dataset_view = dspy.InputField(desc="Dataset name with sample head(5 rows) view") - data_context = dspy.OutputField(desc="Compact JSON describing DuckDB tables, columns, relationships, metrics and use cases") # This is the auto_analyst with planner class auto_analyst(dspy.Module): """Main analyst module that coordinates multiple agents using a planner""" - def __init__(self, agents, retrievers, user_id=None, db_session=None): + def __init__(self, agents, retrievers): # Initialize agent modules and retrievers self.agents = {} self.agent_inputs = {} self.agent_desc = [] - # Load user-enabled template agents if user_id and db_session are provided - if user_id and db_session: - try: - # For planner use, load planner-enabled templates (max 10, prioritized by usage) - template_signatures = load_user_enabled_templates_for_planner_from_db(user_id, db_session) - - # logger.log_message(f"Loaded {template_signatures} templates for planner use", level=logging.INFO) - - for template_name, signature in template_signatures.items(): - # For planner module, load all planner variants (including core planner agents) - # Skip only individual variants, not planner variants - - # Add template agent to agents dict - self.agents[template_name] = dspy.asyncify(dspy.Predict(signature)) - - # Determine if this is a visualization agent based on database category - is_viz_agent = False - try: - from src.db.schemas.models import AgentTemplate - - # Find template record to check category - template_record = db_session.query(AgentTemplate).filter( - AgentTemplate.template_name == template_name - ).first() - - if template_record and template_record.category and template_record.category.lower() == 'visualization': - is_viz_agent = True - else: - # Fallback to name-based detection for legacy templates - is_viz_agent = ('viz' in template_name.lower() or - 'visual' in template_name.lower() or - 'plot' in template_name.lower() or - 'chart' in template_name.lower() or - 'matplotlib' in template_name.lower()) - except Exception as cat_error: - logger.log_message(f"Error checking category for template {template_name}: {str(cat_error)}", level=logging.WARNING) - # Fallback to name-based detection - is_viz_agent = ('viz' in template_name.lower() or - 'visual' in template_name.lower() or - 'plot' in template_name.lower() or - 'chart' in template_name.lower() or - 'matplotlib' in template_name.lower()) - - # Set input fields based on agent type - if is_viz_agent: - self.agent_inputs[template_name] = {'goal', 'dataset', 'styling_index', 'plan_instructions'} - else: - self.agent_inputs[template_name] = {'goal', 'dataset', 'plan_instructions'} - - # Store template agent description - try: - if not template_record: - template_record = db_session.query(AgentTemplate).filter( - AgentTemplate.template_name == template_name - ).first() - - if template_record: - description = f"Template: {template_record.description}" - self.agent_desc.append({template_name: description}) - else: - self.agent_desc.append({template_name: f"Template: {template_name}"}) - except Exception as desc_error: - logger.log_message(f"Error getting description for template {template_name}: {str(desc_error)}", level=logging.WARNING) - self.agent_desc.append({template_name: f"Template: {template_name}"}) - - except Exception as e: - logger.log_message(f"Error loading template agents for user {user_id}: {str(e)}", level=logging.ERROR) - - # Load core planner agents based on user preferences (only planner variants for planner module) - if len(self.agents) == 0 and user_id and db_session: - # try: - # Get user preferences for core planner agents - from src.db.schemas.models import AgentTemplate, UserTemplatePreference - - # For planner module, use planner variants of core agents - core_planner_agent_names = ['planner_preprocessing_agent', 'planner_statistical_analytics_agent', 'planner_sk_learn_agent', 'planner_data_viz_agent'] - - for agent_name in core_planner_agent_names: - # Check if user has enabled this core agent (check both planner and individual preferences) - template = db_session.query(AgentTemplate).filter( - AgentTemplate.template_name == agent_name, - AgentTemplate.is_active == True - ).first() - - if not template: - logger.log_message(f"Core planner agent template '{agent_name}' not found in database", level=logging.WARNING) - continue - - # Check user preference for this planner agent - preference = db_session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - # Core planner agents are enabled by default unless explicitly disabled - is_enabled = preference.is_enabled if preference else True - - if not is_enabled: - continue - - # Skip if already loaded from template_signatures - if agent_name in self.agents: - continue - - # Create dynamic signature for planner agent - signature = create_custom_agent_signature( - template.template_name, - template.description, - template.prompt_template, - template.category - ) - - # Add to agents dict - self.agents[agent_name] = dspy.asyncify(dspy.Predict(signature)) - - # Set input fields based on signature (all planner agents need plan_instructions) - if 'data_viz' in agent_name.lower() or template.category == 'Data Visualization': - self.agent_inputs[agent_name] = {'goal', 'dataset', 'styling_index', 'plan_instructions'} - else: - self.agent_inputs[agent_name] = {'goal', 'dataset', 'plan_instructions'} - - # Get description from database - description = f"Planner: {template.description}" - self.agent_desc.append({agent_name: description}) - logger.log_message(f"Loaded core planner agent: {agent_name}", level=logging.DEBUG) - - # Don't fallback - user must explicitly enable agents - else: - self._load_default_planner_agents_fallback() - # Load standard agents from provided list (legacy support) - - - self.agents['basic_qa_agent'] = dspy.asyncify(dspy.Predict("goal->answer")) - self.agent_inputs['basic_qa_agent'] = {"goal"} - self.agent_desc.append({'basic_qa_agent':"Answers queries unrelated to data & also that include links, poison or attempts to attack the system"}) + for i, a in enumerate(agents): + name = a.__pydantic_core_schema__['schema']['model_name'] + self.agents[name] = dspy.ChainOfThought(a) + self.agent_inputs[name] = {x.strip() for x in str(agents[i].__pydantic_core_schema__['cls']).split('->')[0].split('(')[1].split(',')} + self.agent_desc.append({name: get_agent_description(name)}) # Initialize coordination agents - self.planner = planner_module() - # self.memory_summarize_agent = dspy.ChainOfThought(m.memory_summarize_agent) + self.planner = dspy.ChainOfThought(analytical_planner) + self.refine_goal = dspy.ChainOfThought(goal_refiner_agent) + self.code_combiner_agent = dspy.ChainOfThought(code_combiner_agent) + self.story_teller = dspy.ChainOfThought(story_teller_agent) + self.memory_summarize_agent = dspy.ChainOfThought(m.memory_summarize_agent) # Initialize retrievers - self.dataset = retrievers['dataframe_index'] + self.dataset = retrievers['dataframe_index'].as_retriever(k=1) self.styling_index = retrievers['style_index'].as_retriever(similarity_top_k=1) - # Store user_id for usage tracking - self.user_id = user_id - - - def _load_default_agents_fallback(self): - """Fallback method to load default agents when preference system fails""" - logger.log_message("Loading default agents as fallback for auto_analyst_ind", level=logging.WARNING) - - # Load the 4 core agents from database - core_agent_names = ['preprocessing_agent', 'statistical_analytics_agent', 'sk_learn_agent', 'data_viz_agent'] - - for agent_name in core_agent_names: - # Get the agent signature class - if agent_name == 'preprocessing_agent': - agent_signature = preprocessing_agent - elif agent_name == 'statistical_analytics_agent': - agent_signature = statistical_analytics_agent - elif agent_name == 'sk_learn_agent': - agent_signature = sk_learn_agent - elif agent_name == 'data_viz_agent': - agent_signature = data_viz_agent - - # Add to agents dict - self.agents[agent_name] = dspy.asyncify(dspy.Predict(agent_signature)) - - # Set input fields based on signature - if agent_name == 'data_viz_agent': - self.agent_inputs[agent_name] = {'goal', 'dataset', 'styling_index', 'plan_instructions'} - else: - self.agent_inputs[agent_name] = {'goal', 'dataset', 'plan_instructions'} - - # Get description from database - self.agent_desc.append({agent_name: get_agent_description(agent_name)}) - logger.log_message(f"Added fallback agent: {agent_name}", level=logging.DEBUG) - - def _load_default_planner_agents_fallback(self): - """Fallback method to load default planner agents when preference system fails""" - logger.log_message("Loading default planner agents as fallback for auto_analyst", level=logging.WARNING) - - # For planner module, load the 4 core planner agents - core_planner_agent_names = ['planner_preprocessing_agent', 'planner_statistical_analytics_agent', 'planner_sk_learn_agent', 'planner_data_viz_agent'] - - for agent_name in core_planner_agent_names: - # Skip if already loaded - if agent_name in self.agents: - continue - - # Create a basic signature for the planner agent as fallback - # In production, these should come from the database - if agent_name == 'planner_preprocessing_agent': - base_signature = preprocessing_agent - description = "Planner: Data preprocessing agent for multi-agent pipelines" - elif agent_name == 'planner_statistical_analytics_agent': - base_signature = statistical_analytics_agent - description = "Planner: Statistical analytics agent for multi-agent pipelines" - elif agent_name == 'planner_sk_learn_agent': - base_signature = sk_learn_agent - description = "Planner: Machine learning agent for multi-agent pipelines" - elif agent_name == 'planner_data_viz_agent': - base_signature = data_viz_agent - description = "Planner: Data visualization agent for multi-agent pipelines" - - # Add to agents dict using base signature (fallback mode) - self.agents[agent_name] = dspy.asyncify(dspy.ChainOfThought(base_signature)) - - # Set input fields based on signature - if 'data_viz' in agent_name: - self.agent_inputs[agent_name] = {'goal', 'dataset', 'styling_index', 'plan_instructions'} - else: - self.agent_inputs[agent_name] = {'goal', 'dataset', 'plan_instructions'} - - # Add description - self.agent_desc.append({agent_name: description}) - logger.log_message(f"Added fallback planner agent: {agent_name}", level=logging.DEBUG) + # Initialize thread pool for parallel execution + self.executor = ThreadPoolExecutor(max_workers=min(len(agents) + 2, os.cpu_count() * 2)) - async def _track_agent_usage(self, agent_name): - """Track usage for template agents""" + def execute_agent(self, agent_name, inputs): + """Execute a single agent with given inputs""" try: - # Skip tracking for standard agents and basic_qa_agent (but DO track planner variants) - if agent_name in ['preprocessing_agent', 'statistical_analytics_agent', 'sk_learn_agent', 'data_viz_agent', 'basic_qa_agent']: - return - - # Only track if we have user_id (template agents) - if not self.user_id: - return - - from src.db.init_db import session_factory - from src.db.schemas.models import AgentTemplate, UserTemplatePreference - from datetime import datetime, UTC - - # Create database session - session = session_factory() - try: - # Find the template - template = session.query(AgentTemplate).filter( - AgentTemplate.template_name == agent_name - ).first() - - if not template: - logger.log_message(f"Template '{agent_name}' not found for usage tracking", level=logging.WARNING) - return - - # Find or create user template preference record - preference = session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == self.user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - if not preference: - # Create new preference record (disabled by default) - preference = UserTemplatePreference( - user_id=self.user_id, - template_id=template.template_id, - is_enabled=False, # Disabled by default - usage_count=0, - last_used_at=None, - created_at=datetime.now(UTC), - updated_at=datetime.now(UTC) - ) - session.add(preference) - - # Update usage tracking - preference.usage_count += 1 - preference.last_used_at = datetime.now(UTC) - preference.updated_at = datetime.now(UTC) - session.commit() - - logger.log_message( - f"Tracked usage for template '{agent_name}' (count: {preference.usage_count})", - level=logging.DEBUG - ) - - except Exception as e: - session.rollback() - logger.log_message(f"Error tracking usage for template {agent_name}: {str(e)}", level=logging.ERROR) - finally: - session.close() - + result = self.agents[agent_name.strip()](**inputs) + return agent_name.strip(), dict(result) except Exception as e: - logger.log_message(f"Error in _track_agent_usage for {agent_name}: {str(e)}", level=logging.ERROR) - + return agent_name.strip(), {"error": str(e)} - - async def get_plan(self, query): + def get_plan(self, query): """Get the analysis plan""" dict_ = {} - dict_['dataset'] = self.dataset + dict_['dataset'] = self.dataset.retrieve(query)[0].text dict_['styling_index'] = self.styling_index.retrieve(query)[0].text dict_['goal'] = query dict_['Agent_desc'] = str(self.agent_desc) - - module_return = await self.planner( - goal=dict_['goal'], - dataset=dict_['dataset'], - Agent_desc=dict_['Agent_desc'] - ) - - logger.log_message(f"Module return: {module_return}", level=logging.INFO) - - # Add None check before accessing dictionary keys - if module_return is None: - logger.log_message("Planner returned None, returning error response", level=logging.ERROR) - return { - "plan": "There was an error" + str(dict_) +'\n'+ str(dspy.inspect_history()) +'\n agent_desc_len'+ str(len(self.agent_desc)) + '\n agents_len'+ str(len(self.agents)), - "plan_instructions": {}, - "complexity": "unknown", - "error": "Planner failed to generate a plan" - } - - # Handle different plan formats - plan = module_return['plan'] - logger.log_message(f"Plan from module_return: {plan}, type: {type(plan)}", level=logging.INFO) - - # If plan is a string (agent name), convert to proper format - if isinstance(plan, str): - if 'complexity' in module_return: - complexity = module_return['complexity'] - else: - complexity = 'basic' - - plan_dict = { - 'plan': plan, - 'complexity': complexity - } - - # Add plan_instructions if available - if 'plan_instructions' in module_return: - plan_dict['plan_instructions'] = module_return['plan_instructions'] - else: - plan_dict['plan_instructions'] = {} - else: - # If plan is already a dict, use it directly - plan_dict = dict(plan) if not isinstance(plan, dict) else plan - if 'complexity' in module_return: - complexity = module_return['complexity'] - else: - complexity = 'basic' - plan_dict['complexity'] = complexity - - logger.log_message(f"Final plan dict: {plan_dict}", level=logging.INFO) - - return plan_dict - - # except Exception as e: - # logger.log_message(f"Error in get_plan: {str(e)}", level=logging.ERROR) - # raise + plan = self.planner(goal=dict_['goal'], dataset=dict_['dataset'], Agent_desc=dict_['Agent_desc']) + return dict(plan) async def execute_plan(self, query, plan): """Execute the plan and yield results as they complete""" - dict_ = {} - dict_['dataset'] = self.dataset + dict_['dataset'] = self.dataset.retrieve(query)[0].text dict_['styling_index'] = self.styling_index.retrieve(query)[0].text dict_['hint'] = [] dict_['goal'] = query + import json # Clean and split the plan string into agent names - plan_text = plan.get("plan", "").lower().replace("plan:", "").strip() - logger.log_message(f"Plan text: {plan_text}", level=logging.INFO) - - if "basic_qa_agent" in plan_text: - inputs = dict(goal=query) - - response = await self.agents['basic_qa_agent'](**inputs) - yield 'basic_qa_agent', inputs, response - return - - - plan_list = [] - for agent in [a.strip() for a in plan_text.split("->") if a.strip()]: - if not agent.startswith("planner_"): - agent = "planner_" + agent - plan_list.append(agent) + plan_text = plan.get("plan", "").replace("Plan", "").replace(":", "").strip() + plan_list = [agent.strip() for agent in plan_text.split("->") if agent.strip()] - - - logger.log_message(f"Plan list: {plan_list}", level=logging.INFO) # Parse the attached plan_instructions into a dict raw_instr = plan.get("plan_instructions", {}) if isinstance(raw_instr, str): try: plan_instructions = json.loads(raw_instr) - except Exception as e: - logger.log_message(f"Error parsing plan_instructions JSON: {str(e)}", level=logging.ERROR) + except Exception: plan_instructions = {} elif isinstance(raw_instr, dict): plan_instructions = raw_instr else: plan_instructions = {} - - # Check if we have no valid agents to execute + # If no plan was produced, short-circuit if not plan_list: - if len(plan_text) != 0: - yield "plan_not_formatted_correctly", str(plan_text), {'error': "There was a error in the formatting"} - + yield "plan_not_found", dict(plan), {"error": "No plan found"} return - # Execute agents in sequence - for agent_name in plan_list: - - try: - # Prepare inputs for the agent - inputs = {x: dict_[x] for x in self.agent_inputs[agent_name] if x in dict_} + # Launch each agent in parallel, attaching its own instructions + futures = [] + for idx, agent_name in enumerate(plan_list): + key = agent_name.strip() + # gather input fields except plan_instructions + inputs = { + param: dict_[param] + for param in self.agent_inputs[key] + if param != "plan_instructions" + } + + # attach the specific instructions for this agent with prev/next format + if "plan_instructions" in self.agent_inputs[key]: + # Get current agent instructions + current_instructions = plan_instructions.get(key, {"create": [], "use": [], "instruction": ""}) - # Add plan instructions if available for this agent - if agent_name in plan_instructions: - inputs['plan_instructions'] = plan_instructions[agent_name] - else: - inputs['plan_instructions'] = "" + # Format instructions with your_task first + formatted_instructions = {"your_task": current_instructions} - # logger.log_message(f"Agent inputs for {agent_name}: {inputs}", level=logging.INFO) - - - result = await self.agents[agent_name.strip()](**inputs) + # Add previous agent instructions if available + if idx > 0: + prev_agent = plan_list[idx-1].strip() + prev_instructions = plan_instructions.get(prev_agent, {}).get("instruction", "") + formatted_instructions[f"Previous Agent {prev_agent}"] = prev_instructions - # Track usage for custom agents and templates - await self._track_agent_usage(agent_name.strip()) - # Execute the agent - + # Add next agent instructions if available + if idx < len(plan_list) - 1: + next_agent = plan_list[idx+1].strip() + next_instructions = plan_instructions.get(next_agent, {}).get("instruction", "") + formatted_instructions[f"Next Agent {next_agent}"] = next_instructions - yield agent_name, inputs, result - - except Exception as e: - logger.log_message(f"Error executing agent {agent_name}: {str(e)}", level=logging.ERROR) - yield agent_name, {}, {"error": f"Error executing {agent_name}: {str(e)}"} - return + + inputs["plan_instructions"] = str(formatted_instructions) + logger.log_message(f"Inputs: {inputs}", level=logging.INFO) + future = self.executor.submit(self.execute_agent, agent_name, inputs) + futures.append((agent_name, inputs, future)) - + # Yield results as they complete + completed_results = [] + for agent_name, inputs, future in futures: + try: + name, result = await asyncio.get_event_loop().run_in_executor(None, future.result) + completed_results.append((name, result)) + yield name, inputs, result + except Exception as e: + yield agent_name, inputs, {"error": str(e)} \ No newline at end of file diff --git a/src/agents/deep_agents.py b/src/agents/deep_agents.py deleted file mode 100644 index f23be891364be926e023cde00cdb7c2eb1749783..0000000000000000000000000000000000000000 --- a/src/agents/deep_agents.py +++ /dev/null @@ -1,1110 +0,0 @@ -import asyncio -import ast -import json -import os -import dspy -import numpy as np -import pandas as pd -from dotenv import load_dotenv -from src.utils.logger import Logger -import logging -import datetime -import re -import textwrap - -def clean_print_statements(code_block): - """ - This function cleans up any `print()` statements that might contain unwanted `\n` characters. - It ensures print statements are properly formatted without unnecessary newlines. - """ - # This regex targets print statements, even if they have newlines inside - return re.sub(r'print\((.*?)(\\n.*?)(.*?)\)', r'print(\1\3)', code_block, flags=re.DOTALL) - - -def clean_unicode_chars(text): - """ - Clean Unicode characters that might cause encoding issues. - Replaces common Unicode characters with ASCII equivalents. - """ - if not isinstance(text, str): - return text - - # Replace common Unicode characters with ASCII equivalents - replacements = { - '\u2192': ' -> ', # Right arrow - '\u2190': ' <- ', # Left arrow - '\u2194': ' <-> ', # Left-right arrow - '\u2500': '-', # Box drawing horizontal - '\u2502': '|', # Box drawing vertical - '\u2026': '...', # Ellipsis - '\u2013': '-', # En dash - '\u2014': '-', # Em dash - '\u201c': '"', # Left double quotation mark - '\u201d': '"', # Right double quotation mark - '\u2018': "'", # Left single quotation mark - '\u2019': "'", # Right single quotation mark - } - - for unicode_char, ascii_replacement in replacements.items(): - text = text.replace(unicode_char, ascii_replacement) - - # Remove any remaining non-ASCII characters - text = text.encode('ascii', 'ignore').decode('ascii') - - return text - - -def remove_main_block(code): - # Match the __main__ block - pattern = r'(?m)^if\s+__name__\s*==\s*["\']__main__["\']\s*:\s*\n((?:\s+.*\n?)*)' - - match = re.search(pattern, code) - if match: - main_block = match.group(1) - - # Dedent the code block inside __main__ - dedented_block = textwrap.dedent(main_block) - - # Remove \n from any print statements in the block (also handling multiline print cases) - dedented_block = clean_print_statements(dedented_block) - # Replace the block in the code - cleaned_code = re.sub(pattern, dedented_block, code) - - # Optional: Remove leading newlines if any - cleaned_code = cleaned_code.strip() - - return cleaned_code - return code - - -# Configure Plotly to prevent auto-display -def configure_plotly_no_display(): - """Configure Plotly to prevent automatic browser display""" - try: - import plotly.io as pio - - # Set environment variables to prevent browser opening - os.environ['BROWSER'] = '' - os.environ['PLOTLY_RENDERER'] = 'json' - - # Configure Plotly renderers - pio.renderers.default = 'json' - pio.templates.default = 'plotly_white' - - # Disable Kaleido auto-display if available - try: - import plotly.graph_objects as go - # Configure figure defaults to not auto-display - go.Figure.show = lambda self, *args, **kwargs: None - except ImportError: - pass - - except ImportError: - print("Warning: Plotly not available for configuration") - -# Call the configuration function immediately -configure_plotly_no_display() - -logger = Logger("deep_agents", see_time=True, console_log=True, level=logging.DEBUG) -load_dotenv() - -class deep_questions(dspy.Signature): - """ -You are a data analysis assistant. -Your role is to take a user's high-level analytical goal and generate a set of deep, targeted follow-up questions. These questions should guide an analyst toward a more thorough understanding of the goal by encouraging exploration, segmentation, and causal reasoning. -Instructions: -- Generate up to 5 insightful, data-relevant questions. -- Use the dataset structure to tailor your questions (e.g., look at the available columns, data types, and what kind of information they can reveal). -- The questions should help the user decompose their analytic goal and explore it from multiple angles (e.g., time trends, customer segments, usage behavior, external factors, feedback). -- Each question should be specific enough to guide actionable analysis or investigation. -- Use a clear and concise style, but maintain depth. -Inputs: -- goal: The user's analytical goal or main question they want to explore -- dataset_info: A description of the dataset the user is querying, including: - - What the dataset represents - - Key columns and their data types -Output: -- deep_questions: A list of up to 5 specific, data-driven questions that support the analytic goal ---- -Example: -Analytical Goal: -Understand why churn has been rising -Dataset Info: -Customer Retention Dataset tracking subscription activity over time. -Columns: -- customer_id (string) -- join_date (date) -- churn_date (date, nullable) -- is_churned (boolean) -- plan_type (string: 'basic', 'premium', 'enterprise') -- region (string) -- last_login_date (date) -- avg_weekly_logins (float) -- support_tickets_last_30d (int) -- satisfaction_score (float, 0–10 scale) -Decomposed Questions: -1. How has the churn rate changed month-over-month, and during which periods was the increase most pronounced? -2. Are specific plan types or regions showing a higher churn rate relative to others? -3. What is the average satisfaction score and support ticket count among churned users compared to retained users? -4. Do churned users exhibit different login behavior (e.g., avg_weekly_logins) in the weeks leading up to their churn date? -5. What is the tenure distribution (time from join_date to churn_date) among churned customers, and are short-tenure users more likely to churn? - """ - goal = dspy.InputField(desc="User analytical goal — what main insight or question they want to answer") - dataset_info = dspy.InputField(desc="A description of the dataset: what it represents, and the main columns with data types") - deep_questions:str = dspy.OutputField(desc="A list of up to five questions that help deeply explore the analytical goal using the dataset") - -class deep_synthesizer(dspy.Signature): - """ -You are a data analysis synthesis expert. -Your job is to take the outputs from a multi-agent data analytics system - including the original user query, the code summaries from each agent, and the actual printed results from running those code blocks - and synthesize them into a comprehensive, well-structured final report. -This report should: -- Explain what steps were taken and why (based on the query) -- Summarize the code logic used by each agent, without including raw code -- Highlight key findings and results from the code outputs -- Offer clear, actionable insights tied back to the user's original question -- Be structured, readable, and suitable for decision-makers or analysts -Instructions: -- Begin with a brief restatement of the original query and what it aimed to solve -- Organize your report step-by-step or by analytical theme (e.g., segmentation, trend analysis, etc.) -- For each part, summarize what was analyzed, how (based on code summaries), and what the result was (based on printed output) -- End with a final set of synthesized conclusions and potential next steps or recommendations -Inputs: -- query: The user's original analytical question or goal -- summaries: A list of natural language descriptions of what each agent's code did -- print_outputs: A list of printed outputs (results) from running each agent's code -Output: -- synthesized_report: A structured and readable report that ties all parts together, grounded in the code logic and results -Example use: -You are not just summarizing outputs - you're telling a story that answers the user's query using real data. - """ - - query = dspy.InputField(desc="The original user query or analytical goal") - summaries = dspy.InputField(desc="List of code summaries - each describing what a particular agent's code did") - print_outputs = dspy.InputField(desc="List of print outputs - the actual data insights generated by the code") - synthesized_report = dspy.OutputField(desc="The final, structured report that synthesizes all the information into clear insights") - -def clean_and_store_code(code, session_datasets=None): - """ - Cleans and stores code execution results in a standardized format. - - Args: - code (str): Raw code text to execute - session_datasets (dict): Dictionary of datasets from session state - - Returns: - dict: Execution results containing printed_output, plotly_figs, and error info - """ - import io - import sys - import re - import plotly.express as px - import plotly.graph_objects as go - from plotly.subplots import make_subplots - import plotly.io as pio - - # Make all session datasets available globally if provided - if session_datasets is not None: - for dataset_name, dataset_df in session_datasets.items(): - globals()[dataset_name] = dataset_df - - # Initialize output containers - output_dict = { - 'exec_result': None, - 'printed_output': '', - 'plotly_figs': [], - 'error': None - } - - try: - # Clean the code - cleaned_code = code.strip() - - cleaned_code = cleaned_code.replace('```python', '').replace('```', '') - - - # Fix try statement syntax - cleaned_code = cleaned_code.replace('try\n', 'try:\n') - - # Remove code patterns that would make the code unrunnable - invalid_patterns = [ - '```', # Code block markers - '\\n', # Raw newlines - '\\t', # Raw tabs - '\\r', # Raw carriage returns - ] - - for pattern in invalid_patterns: - if pattern in cleaned_code: - cleaned_code = cleaned_code.replace(pattern, '') - - - # Remove reading the csv file if it's already in the context - cleaned_code = re.sub(r"df\s*=\s*pd\.read_csv\([\"\'].*?[\"\']\).*?(\n|$)", '', cleaned_code) - - # Only match assignments at top level (not indented) - # 1. Remove 'df = pd.DataFrame()' if it's at the top level - cleaned_code = re.sub( - r"^df\s*=\s*pd\.DataFrame\(\s*\)\s*(#.*)?$", - '', - cleaned_code, - flags=re.MULTILINE - ) - cleaned_code = re.sub(r"plt\.show\(\).*?(\n|$)", '', cleaned_code) - # Remove all .show() method calls more comprehensively - cleaned_code = re.sub(r'\b\w*\.show\(\)', '', cleaned_code) - cleaned_code = re.sub(r'^\s*\w*fig\w*\.show\(\)\s*;?\s*$', '', cleaned_code, flags=re.MULTILINE) - - # Additional patterns to catch more .show() variations - cleaned_code = re.sub(r'\.show\(\s*\)', '', cleaned_code) # .show() with optional spaces - cleaned_code = re.sub(r'\.show\(\s*renderer\s*=\s*[\'"][^\'\"]*[\'"]\s*\)', '', cleaned_code) # .show(renderer='...') - cleaned_code = re.sub(r'plotly_figs\[\d+\]\.show\(\)', '', cleaned_code) # plotly_figs[0].show() - - # More comprehensive patterns - cleaned_code = re.sub(r'\.show\([^)]*\)', '', cleaned_code) # .show(any_args) - cleaned_code = re.sub(r'fig\w*\.show\(\s*[^)]*\s*\)', '', cleaned_code) # fig*.show(any_args) - cleaned_code = re.sub(r'\w+_fig\w*\.show\(\s*[^)]*\s*\)', '', cleaned_code) # *_fig*.show(any_args) - - cleaned_code = remove_main_block(cleaned_code) - - # Clean Unicode characters that might cause encoding issues - cleaned_code = clean_unicode_chars(cleaned_code) - - # Capture printed output - old_stdout = sys.stdout - captured_output = io.StringIO() - sys.stdout = captured_output - - # Create execution environment with common imports and session data - exec_globals = { - '__builtins__': __builtins__, - 'pd': __import__('pandas'), - 'np': __import__('numpy'), - 'px': px, - 'go': go, - 'make_subplots': make_subplots, - 'plotly_figs': [], - 'print': print, - } - - # Add session DataFrame if available - if session_datasets is not None: - for dataset_name, dataset_df in session_datasets.items(): - exec_globals[dataset_name] = dataset_df - - # Add other common libraries that might be needed - try: - exec_globals['sm'] = __import__('statsmodels.api', fromlist=['']) - exec_globals['train_test_split'] = __import__('sklearn.model_selection', fromlist=['train_test_split']).train_test_split - exec_globals['LinearRegression'] = __import__('sklearn.linear_model', fromlist=['LinearRegression']).LinearRegression - exec_globals['mean_absolute_error'] = __import__('sklearn.metrics', fromlist=['mean_absolute_error']).mean_absolute_error - exec_globals['r2_score'] = __import__('sklearn.metrics', fromlist=['r2_score']).r2_score - exec_globals['LabelEncoder'] = __import__('sklearn.preprocessing', fromlist=['LabelEncoder']).LabelEncoder - exec_globals['warnings'] = __import__('warnings') - except ImportError as e: - print(f"Warning: Could not import some optional libraries: {e}") - - # Execute the code - exec(cleaned_code, exec_globals) - - # Restore stdout - sys.stdout = old_stdout - - # Get the captured output - printed_output = captured_output.getvalue() - output_dict['printed_output'] = printed_output - # Extract plotly figures from the execution environment - if 'plotly_figs' in exec_globals: - plotly_figs = exec_globals['plotly_figs'] - if isinstance(plotly_figs, list): - output_dict['plotly_figs'] = plotly_figs - else: - output_dict['plotly_figs'] = [plotly_figs] if plotly_figs else [] - - # Also check for any figure variables that might have been created - for var_name, var_value in exec_globals.items(): - if hasattr(var_value, 'to_json') and hasattr(var_value, 'show'): - # This looks like a Plotly figure - if var_value not in output_dict['plotly_figs']: - output_dict['plotly_figs'].append(var_value) - - except Exception as e: - # Restore stdout in case of error - sys.stdout = old_stdout - error_msg = str(e) - output_dict['error'] = error_msg - output_dict['printed_output'] = f"Error executing code: {error_msg}" - print(f"Code execution error: {error_msg}") - - return output_dict - -def score_code(args, code, datasets=None): - """ - Cleans and stores code execution results in a standardized format. - Safely handles execution errors and returns clean output even if execution fails. - Ensures plotly figures are properly created and captured. - - Args: - args: Arguments (unused but required for dspy.Refine) - code: Code object with combined_code attribute - datasets: Dictionary of datasets from session state (optional) - - Returns: - int: Score (0=error, 1=success, 2=success with plots) - """ - - code_text = code.combined_code - try: - # Fix try statement syntax - code_text = code_text.replace('try\n', 'try:\n') - code_text = code_text.replace('```python', '').replace('```', '') - - - # Remove code patterns that would make the code unrunnable - invalid_patterns = [ - '```', '\\n', '\\t', '\\r' - ] - - for pattern in invalid_patterns: - if pattern in code_text: - code_text = code_text.replace(pattern, '') - - cleaned_code = re.sub(r"plt\.show\(\).*?(\n|$)", '', code_text) - # Remove all .show() method calls more comprehensively - cleaned_code = re.sub(r'\b\w*\.show\(\)', '', cleaned_code) - cleaned_code = re.sub(r'^\s*\w*fig\w*\.show\(\)\s*;?\s*$', '', cleaned_code, flags=re.MULTILINE) - - # Additional patterns to catch more .show() variations - cleaned_code = re.sub(r'\.show\(\s*\)', '', cleaned_code) # .show() with optional spaces - cleaned_code = re.sub(r'\.show\(\s*renderer\s*=\s*[\'"][^\'\"]*[\'"]\s*\)', '', cleaned_code) # .show(renderer='...') - cleaned_code = re.sub(r'plotly_figs\[\d+\]\.show\(\)', '', cleaned_code) # plotly_figs[0].show() - - # More comprehensive patterns - cleaned_code = re.sub(r'\.show\([^)]*\)', '', cleaned_code) # .show(any_args) - cleaned_code = re.sub(r'fig\w*\.show\(\s*[^)]*\s*\)', '', cleaned_code) # fig*.show(any_args) - cleaned_code = re.sub(r'\w+_fig\w*\.show\(\s*[^)]*\s*\)', '', cleaned_code) # *_fig*.show(any_args) - - cleaned_code = remove_main_block(cleaned_code) - - # Capture stdout using StringIO - from io import StringIO - import sys - import plotly.graph_objects as go - import pandas as pd - import numpy as np - - stdout_capture = StringIO() - original_stdout = sys.stdout - sys.stdout = stdout_capture - - # Execute code in a new namespace with datasets available - local_vars = {} - - # Add datasets to the execution context if provided - if datasets: - local_vars.update(datasets) - - # Add common imports to the execution context - local_vars.update({ - 'pd': pd, - 'np': np, - 'go': go, - 'plt': __import__('matplotlib.pyplot'), - 'sns': __import__('seaborn'), - }) - - exec(cleaned_code, globals(), local_vars) - - # Capture any plotly figures from local namespace - plotly_figs = [] - for var_name, var in local_vars.items(): - if isinstance(var, go.Figure): - if not var.layout.title: - var.update_layout(title=f"Figure {len(plotly_figs) + 1}") - if not var.layout.template: - var.update_layout(template="plotly_white") - plotly_figs.append(var) - elif isinstance(var, (list, tuple)): - for item in var: - if isinstance(item, go.Figure): - if not item.layout.title: - item.update_layout(title=f"Figure {len(plotly_figs) + 1}") - if not item.layout.template: - item.update_layout(template="plotly_white") - plotly_figs.append(item) - - # Restore stdout and get captured output - sys.stdout = original_stdout - captured_output = stdout_capture.getvalue() - stdout_capture.close() - - # Calculate score based on execution and plot generation - score = 2 if plotly_figs else 1 - - return score - - except Exception as e: - # Restore stdout in case of error - if 'stdout_capture' in locals(): - sys.stdout = original_stdout - stdout_capture.close() - - return 0 - - -class deep_planner(dspy.Signature): - """ - You are an advanced multi-question planning agent. Your task is to generate the most optimized and minimal plan - to answer up to 5 analytical questions using available agents. - Your responsibilities: - 1. Feasibility: Verify that the goal is achievable using the provided datasets and agent descriptions. - 2. Optimization: - - Batch up to 2 similar questions per agent call. - - Reuse outputs across questions wherever possible. - - Avoid unnecessary agents or redundant processing. - - Minimize total agent calls while preserving correctness. - 3. Clarity: - - Define clear variable usage (create/use). - - Specify concise step-by-step instructions per agent. - - Use dependency arrows (->) to indicate required agent outputs used by others. - Inputs: - - deep_questions: A list of up to 5 deep analytical questions (e.g., ["q1", "q2", ..., "q5"]) - - dataset: The available dataset(s) in memory or context - - agents_desc: Dictionary containing each agent's name and its capabilities or descriptions - Outputs: - - plan_instructions: Detailed per-agent variable flow and functionality in the format: - { - "agent_x": { - "create": ["cleaned_data: DataFrame - cleaned version of the input dataset"], - "use": ["df: DataFrame - raw input dataset"], - "instruction": "Clean the dataset by handling null values and standardizing formats." - }, - "agent_y": { - "create": ["analysis_results: dict - results of correlation analysis"], - "use": ["cleaned_data: DataFrame - output from @agent_x"], - "instruction": "Perform correlation analysis to identify strong predictors." - } - } - Output Goal: - Generate a small, clean, optimized execution plan using minimal agent calls, reusable outputs, and well-structured dependencies. - USE THE EXACT NAME OF THE AGENTS IN THE INSTRUCTIONS - """ - - deep_questions = dspy.InputField(desc="List of up to 5 deep analytical questions to answer") - dataset = dspy.InputField(desc="Available datasets, use 'df' as the working dataset") - agents_desc = dspy.InputField(desc="Descriptions of available agents and their functions") - plan_instructions = dspy.OutputField(desc="Variable-level instructions for each agent used in the plan") - -class deep_plan_fixer(dspy.Signature): - """ - You are a plan instruction fixer agent. Your task is to take potentially malformed plan instructions - and convert them into a properly structured dictionary format that can be safely evaluated. - Your responsibilities: - 1. Parse and validate the input plan instructions - 2. Convert the instructions into a proper dictionary format - 3. Ensure all agent instructions follow the required structure: - { - "@agent_name": { - "create": ["variable: type - description"], - "use": ["variable: type - description"], - "instruction": "clear instruction text" - } - } - 4. Handle any malformed or missing components - 5. Return a properly formatted dictionary string that can be safely evaluated - Inputs: - - plan_instructions: The potentially malformed plan instructions to fix - Outputs: - - fixed_plan: A properly formatted dictionary string that can be safely evaluated - """ - - plan_instructions = dspy.InputField(desc="The potentially malformed plan instructions to fix") - fixed_plan = dspy.OutputField(desc="Properly formatted dictionary string that can be safely evaluated") - -class final_conclusion(dspy.Signature): - """ -You are a high-level analytics reasoning engine. -Your task is to take multiple synthesized analytical results (each answering part of the original query) and produce a cohesive final conclusion that directly addresses the user's original question. -This is not just a summary — it's a judgment. Use evidence from the synthesized findings to: -- Answer the original question with clarity -- Highlight the most important insights -- Offer any causal reasoning or patterns discovered -- Suggest next steps or strategic recommendations where appropriate -Instructions: -- Focus on relevance to the original query -- Do not just repeat what the synthesized sections say — instead, infer, interpret, and connect dots -- Prioritize clarity and insight over detail -- End with a brief "Next Steps" section if applicable -Inputs: -- query: The original user question or goal -- synthesized_sections: A list of synthesized result sections from the deep_synthesizer step (each covering part of the analysis) -Output: -- final_summary: A cohesive final conclusion that addresses the query, draws insight, and offers high-level guidance ---- -Example Output Structure: -**Conclusion** -Summarize the overall answer to the user's question, using the most compelling evidence across the synthesized sections. -**Key Takeaways** -- Bullet 1 -- Bullet 2 -- Bullet 3 -**Recommended Next Steps** -(Optional based on context) - """ - - query = dspy.InputField(desc="The user's original query or analytical goal") - synthesized_sections = dspy.InputField(desc="List of synthesized outputs — each one corresponding to a sub-part of the analysis") - final_conclusion = dspy.OutputField(desc="A cohesive, conclusive answer that addresses the query and integrates key insights") - - - - -class deep_code_synthesizer(dspy.Signature): - """ -You are a code synthesis and optimization engine that combines and fixes code from multiple analytical agents. -Your task is to take code outputs from preprocessing, statistical analysis, machine learning, and visualization agents, then: -- Combine them into a single, coherent analysis pipeline -- Fix any errors or inconsistencies between agent outputs -- Ensure proper data flow between steps -- Optimize the combined code for efficiency -- Add necessary imports and dependencies -- Handle any data type mismatches or conversion issues -- Validate and normalize data types between agent outputs (e.g., ensure DataFrame operations maintain DataFrame type) -- Convert between common data structures (lists, dicts, DataFrames) as needed -- Add type hints and validation checks -- Ensure consistent variable naming across agents -- Ensure all visualizations use Plotly exclusively -- Create comprehensive visualizations that show all important variables and relationships -- Store all Plotly figures in a list for later use in the report -Instructions: -- Review each agent's code for correctness and completeness -- Ensure variables are properly passed between steps with consistent types -- Fix any syntax errors or logical issues -- Add error handling and type validation where needed -- Optimize code structure and performance -- Maintain consistent coding style -- Add clear comments explaining the analysis flow -- Add data type conversion functions where needed -- Validate input/output types between agent steps -- Handle edge cases where agents might return different data structures -- Convert any non-Plotly visualizations to Plotly format -- Ensure all important variables are visualized appropriately -- Store all Plotly figures in a list called plotly_figs -- Include appropriate titles, labels, and legends for all visualizations -- Use consistent styling across all Plotly visualizations -- DONOT COMMENT OUT ANYTHING AS THE CODE SHOULD RUN & SHOW OUTPUTS -- THE DATASET IS ALREADY LOADED, DON'T CREATE FAKE DATA. 'df' is always loaded -Inputs: -- deep_questions- The five deep questions this system is answering -- dataset_info - Information about the dataset structure and types -- planner_instructions - the plan according to the planner, ensure that the final code makes everything coherent -- code - List of all agent code -Output: -- combined_code: - A single, optimized Python script that combines all analysis steps with proper type handling and Plotly visualizations -""" - deep_questions = dspy.InputField(desc="The five deep questions this system is answering") - dataset_info = dspy.InputField(desc="Information about the dataset") - planner_instructions = dspy.InputField(desc="The planner instructions for each") - code = dspy.InputField(desc="The code generated by all agents") - combined_code = dspy.OutputField(desc="A single, optimized Python script that combines all analysis steps") - -class deep_code_fix(dspy.Signature): - """ - You are a code debugging and fixing agent that analyzes and repairs code errors. - - Your task is to: - - Analyze error messages and identify root causes - - Fix syntax errors, logical issues, and runtime problems - - Ensure proper data type handling and conversions - - Add appropriate error handling and validation - - Maintain code style and documentation - - Preserve the original analysis intent - - Instructions: - - Carefully analyze the error message and stack trace - - Identify the specific line(s) causing the error - - Determine if the issue is syntax, logic, or runtime related - - Fix the code while maintaining its original purpose - - Add appropriate error handling if needed - - Ensure the fix doesn't introduce new issues - - Document the changes made - - Inputs: - - code: The code that generated the error - - error: The error message and stack trace - - Output: - - fixed_code: The repaired code with error handling - - fix_explanation: Explanation of what was fixed and why - """ - code = dspy.InputField(desc="The code that generated the error") - error = dspy.InputField(desc="The error message and stack trace") - fixed_code = dspy.OutputField(desc="The repaired code with error handling") - fix_explanation = dspy.OutputField(desc="Explanation of what was fixed and why") - - -chart_instructions = """ -Chart Styling Guidelines: -1. General Styling: - - Use a clean, professional color palette (e.g., Tableau, ColorBrewer) - - Include clear titles and axis labels - - Add appropriate legends - - Use consistent font sizes and styles - - Include grid lines where helpful - - Add hover information for interactive plots -2. Specific Chart Types: - - Bar Charts: - * Use horizontal bars for many categories - * Sort bars by value when appropriate - * Use consistent bar widths - * Add value labels on bars - - - Line Charts: - * Use distinct line styles/colors - * Add markers at data points - * Include trend lines when relevant - * Show confidence intervals if applicable - - - Scatter Plots: - * Use appropriate marker sizes - * Add regression lines when needed - * Use color to show additional dimensions - * Include density contours for large datasets - - - Heatmaps: - * Use diverging color schemes for correlation - * Include value annotations - * Sort rows/columns by similarity - * Add clear color scale legend -3. Data Visualization Best Practices: - - Start axes at zero when appropriate - - Use log scales for wide-ranging data - - Include reference lines/benchmarks - - Add annotations for important points - - Show uncertainty where relevant - - Use consistent color encoding - - Include data source and timestamp - - Add clear figure captions -4. Interactive Features: - - Enable zooming and panning - - Add tooltips with detailed information - - Include download options - - Allow toggling of data series - - Enable cross-filtering between charts -5. Accessibility: - - Use colorblind-friendly palettes - - Include alt text for all visualizations - - Ensure sufficient contrast - - Make interactive elements keyboard accessible - - Provide text alternatives for key insights -""" - - - - -class deep_analysis_module(dspy.Module): - def __init__(self,agents, agents_desc): - # logger.log_message(f"Initializing deep_analysis_module with {agents} agents: {list(agents.keys())}", level=logging.INFO) - - # self.agent = agents - self.agents = {} - for key in agents.keys(): - self.agents[key] = dspy.asyncify(dspy.Predict(agents[key])) - - # for sig in self.agen - # Make all dspy operations async using asyncify - self.deep_questions = dspy.asyncify(dspy.Predict(deep_questions)) - self.deep_planner = dspy.asyncify(dspy.Predict(deep_planner)) - self.deep_synthesizer = dspy.asyncify(dspy.Predict(deep_synthesizer)) - # Keep both asyncified and non-asyncified versions for code synthesizer - self.deep_code_synthesizer_sync = dspy.Predict(deep_code_synthesizer) # For dspy.Refine - self.deep_code_synthesizer = dspy.asyncify(dspy.Predict(deep_code_synthesizer)) # For async use - self.deep_plan_fixer = dspy.asyncify(dspy.Predict(deep_plan_fixer)) - self.deep_code_fixer = dspy.asyncify(dspy.Predict(deep_code_fix)) - self.styling_instructions = chart_instructions - self.agents_desc = agents_desc - self.final_conclusion = dspy.asyncify(dspy.ChainOfThought(final_conclusion)) - - # logger.log_message(f"Deep analysis module initialized successfully with agents: {list(self.agents.keys())}", level=logging.INFO) - - async def execute_deep_analysis_streaming(self, goal, dataset_info, session_datasets=None): - """ - Execute deep analysis with streaming progress updates. - This is an async generator that yields progress updates incrementally. - """ - logger.log_message(f"🔍 DEEP ANALYSIS STREAMING START - goal: {goal[:100]}...", level=logging.DEBUG) - logger.log_message(f"🔍 DEEP ANALYSIS STREAMING START - dataset_info type: {type(dataset_info)}, length: {len(dataset_info) if isinstance(dataset_info, str) else 'N/A'}", level=logging.DEBUG) - logger.log_message(f"🔍 DEEP ANALYSIS STREAMING START - session_datasets type: {type(session_datasets)}, keys: {list(session_datasets.keys()) if session_datasets else 'None'}", level=logging.DEBUG) - - # Make all session datasets available globally for code execution - if session_datasets is not None: - for dataset_name, dataset_df in session_datasets.items(): - globals()[dataset_name] = dataset_df - logger.log_message(f"🔍 MADE DATASET AVAILABLE GLOBALLY - {dataset_name}: shape {dataset_df.shape}", level=logging.DEBUG) - - try: - # Step 1: Generate deep questions (20% progress) - yield { - "step": "questions", - "status": "processing", - "message": "Generating analytical questions...", - "progress": 10 - } - - logger.log_message(f"🔍 CALLING DEEP_QUESTIONS - dataset_info type: {type(dataset_info)}, length: {len(dataset_info) if isinstance(dataset_info, str) else 'N/A'}", level=logging.DEBUG) - questions = await self.deep_questions(goal=goal, dataset_info=dataset_info) - logger.log_message("Questions generated") - - yield { - "step": "questions", - "status": "completed", - "content": questions.deep_questions, - "progress": 20 - } - - # Step 2: Create analysis plan (40% progress) - yield { - "step": "planning", - "status": "processing", - "message": "Creating analysis plan...", - "progress": 25 - } - - question_list = [q.strip() for q in questions.deep_questions.split('\n') if q.strip()] - logger.log_message(f" CALLING DEEP_PLANNER - dataset_info type: {type(dataset_info)}, length: {len(dataset_info) if isinstance(dataset_info, str) else 'N/A'}", level=logging.DEBUG) - deep_plan = await self.deep_planner( - deep_questions=questions.deep_questions, - dataset=dataset_info, - agents_desc=str(self.agents_desc) - ) - logger.log_message("Plan created") - - # Parse plan instructions - try: - plan_instructions = ast.literal_eval(deep_plan.plan_instructions) - if not isinstance(plan_instructions, dict): - plan_instructions = json.loads(deep_plan.plan_instructions) - keys = [key for key in plan_instructions.keys()] - - if not all(key in self.agents for key in keys): - raise ValueError(f"Invalid agent key(s) in plan instructions. Available agents: {list(self.agents.keys())}") - logger.log_message(f"Plan instructions: {plan_instructions}", logging.INFO) - logger.log_message(f"Keys: {keys}", logging.INFO) - except (ValueError, SyntaxError, json.JSONDecodeError) as e: - try: - deep_plan = await self.deep_plan_fixer(plan_instructions=deep_plan.plan_instructions) - plan_instructions = ast.literal_eval(deep_plan.fixed_plan) - if not isinstance(plan_instructions, dict): - plan_instructions = json.loads(deep_plan.fixed_plan) - keys = [key for key in plan_instructions.keys()] - except (ValueError, SyntaxError, json.JSONDecodeError) as e: - logger.log_message(f"Error parsing plan instructions: {e}", logging.ERROR) - raise e - - logger.log_message("Instructions parsed") - - yield { - "step": "planning", - "status": "completed", - "content": deep_plan.plan_instructions, - "progress": 40 - } - - # Step 3: Execute agent tasks (60% progress) - yield { - "step": "agent_execution", - "status": "processing", - "message": "Executing analysis agents...", - "progress": 45 - } - - logger.log_message(f"🔍 CREATING DSPY EXAMPLES - dataset_info type: {type(dataset_info)}, length: {len(dataset_info) if isinstance(dataset_info, str) else 'N/A'}", level=logging.DEBUG) - - - # tasks = [self.agents[key](**q) for q, key in zip(queries, keys)] - tasks = [] - for key in keys: - logger.log_message(f"Preparing agent task for key: {key}", level=logging.DEBUG) - logger.log_message(f"Plan instructions for {key}: {plan_instructions[key]}", level=logging.DEBUG) - if "data_viz" in key or "viz" in key.lower() or "visual" in key.lower() or "plot" in key.lower() or "chart" in key.lower(): - logger.log_message(f"Agent {key} identified as visualization agent. Adding styling_index.", level=logging.DEBUG) - task = self.agents[key]( - goal=questions.deep_questions, - dataset=dataset_info, - plan_instructions=str(plan_instructions[key]), - styling_index=self.styling_instructions - ) - logger.log_message(f"Visualization agent task for {key} created.", level=logging.DEBUG) - tasks.append(task) - else: - logger.log_message(f"Agent {key} identified as non-visualization agent.", level=logging.DEBUG) - task = self.agents[key]( - goal=questions.deep_questions, - dataset=dataset_info, - plan_instructions=str(plan_instructions[key]) - ) - logger.log_message(f"Non-visualization agent task for {key} created.", level=logging.DEBUG) - tasks.append(task) - - - - - - - logger.log_message("Passed that stage") - # DEBUG: Log what parameters each agent will receive - - # Await all tasks to complete - summaries = [] - codes = [] - logger.log_message("Tasks started") - - completed_tasks = 0 - for task in asyncio.as_completed(tasks): - result = await task - summaries.append(result.summary) - codes.append(result.code) - completed_tasks += 1 - - # Update progress for each completed agent - agent_progress = 45 + (completed_tasks / len(tasks)) * 15 # 45% to 60% - yield { - "step": "agent_execution", - "status": "processing", - "message": f"Completed {completed_tasks}/{len(tasks)} analysis agents...", - "progress": int(agent_progress) - } - logger.log_message(f"Done with agent {completed_tasks}/{len(tasks)}") - - yield { - "step": "agent_execution", - "status": "completed", - "message": "All analysis agents completed", - "progress": 60 - } - - # Step 4: Code synthesis (80% progress) - yield { - "step": "code_synthesis", - "status": "processing", - "message": "Analyzing code...", - "progress": 65 - } - - # Safely extract code from agent outputs - code = [] - for c in codes: - try: - cleaned_code = remove_main_block(c) - if "```python" in cleaned_code: - parts = cleaned_code.split("```python") - if len(parts) > 1: - extracted = parts[1].split("```")[0] if "```" in parts[1] else parts[1] - code.append(extracted.replace('try\n','try:\n')) - else: - code.append(cleaned_code.replace('try\n','try:\n')) - else: - code.append(cleaned_code.replace('try\n','try:\n')) - except Exception as e: - logger.log_message(f"Warning: Error processing code block: {e}", logging.WARNING) - code.append(c.replace('try\n','try:\n')) - - # Create deep coder without asyncify to avoid source inspection issues - def create_score_code_with_datasets(datasets): - """ - Creates a score function that includes the session datasets. - - Args: - datasets: Dictionary of datasets from session_state['datasets'] - - Returns: - A reward function compatible with dspy.Refine - """ - def score_code_with_datasets(args, pred): - return score_code(args, pred, datasets) # FIXED: Use positional argument - - return score_code_with_datasets - - # Then in your deep analysis method: - # Create score function with datasets - score_fn = create_score_code_with_datasets(self.datasets) - deep_coder = dspy.Refine(module=self.deep_code_synthesizer_sync, N=5, reward_fn=score_fn, threshold=1.0, fail_count=10) - - # Check if we have valid API key - anthropic_key = os.environ.get('ANTHROPIC_API_KEY') - if not anthropic_key: - raise ValueError("ANTHROPIC_API_KEY environment variable is not set") - - try: - # Create the LM instance that will be used - thread_lm = dspy.LM("anthropic/claude-sonnet-4-6", api_key=anthropic_key, max_tokens=17000) - - logger.log_message("Starting code generation...") - start_time = datetime.datetime.now() - logger.log_message(f"Code generation started at: {start_time.strftime('%Y-%m-%d %H:%M:%S')}") - - # Define the blocking function to run in thread - def run_deep_coder(): - with dspy.context(lm=thread_lm): - return deep_coder( - deep_questions=str(questions.deep_questions), - dataset_info=dataset_info, - planner_instructions=str(plan_instructions), - code=str(code) - ) - - # Use asyncio.to_thread for better async integration - deep_code = await asyncio.to_thread(run_deep_coder) - - logger.log_message(f"Code generation completed at: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") - except Exception as e: - logger.log_message(f"Error during code generation: {str(e)}", logging.ERROR) - raise e - - code = deep_code.combined_code - code = code.replace('```python', '').replace('```', '') - - # Clean Unicode characters that might cause encoding issues - code = clean_unicode_chars(code) - - yield { - "step": "code_synthesis", - "status": "completed", - "message": "Code synthesis completed", - "progress": 80 - } - - # Step 5: Execute code (85% progress) - yield { - "step": "code_execution", - "status": "processing", - "message": "Executing code...", - "progress": 82 - } - - # Execute the code with error handling and session DataFrame - try: - # Run code execution in thread pool to avoid blocking - import concurrent.futures - with concurrent.futures.ThreadPoolExecutor() as executor: - future = executor.submit(clean_and_store_code, code, session_datasets) - output = future.result(timeout=300) # 5 minute timeout - - logger.log_message(f"Deep Code executed") - - if output.get('error'): - logger.log_message(f"Warning: Code execution had errors: {output['error']}", logging.ERROR) - - print_outputs = [output['printed_output']] - plotly_figs = [output['plotly_figs']] - - except Exception as e: - logger.log_message(f"Error during code execution: {str(e)}", logging.ERROR) - output = { - 'exec_result': None, - 'printed_output': f"Code execution failed: {str(e)}", - 'plotly_figs': [], - 'error': str(e) - } - print_outputs = [output['printed_output']] - plotly_figs = [output['plotly_figs']] - - yield { - "step": "code_execution", - "status": "completed", - "message": "Code execution completed", - "progress": 85 - } - - # Step 6: Synthesis (90% progress) - yield { - "step": "synthesis", - "status": "processing", - "message": "Synthesizing results...", - "progress": 87 - } - - synthesis = [] - try: - synthesis_result = await self.deep_synthesizer( - query=goal, - summaries=str(summaries), - print_outputs=str(output['printed_output']) - ) - synthesis.append(synthesis_result) - except Exception as e: - logger.log_message(f"Error during synthesis: {str(e)}", logging.ERROR) - synthesis.append(type('obj', (object,), {'synthesized_report': f"Synthesis failed: {str(e)}"})()) - - logger.log_message("Synthesis done") - - yield { - "step": "synthesis", - "status": "completed", - "message": "Synthesis completed", - "progress": 90 - } - - # Step 7: Final conclusion (100% progress) - yield { - "step": "conclusion", - "status": "processing", - "message": "Generating final conclusion...", - "progress": 95 - } - - try: - final_conclusion = await self.final_conclusion( - query=goal, - synthesized_sections=str([s.synthesized_report for s in synthesis]) - ) - except Exception as e: - logger.log_message(f"Error during final conclusion: {str(e)}", logging.ERROR) - final_conclusion = type('obj', (object,), {'final_conclusion': f"Final conclusion failed: {str(e)}"})() - - logger.log_message("Conclusion Made") - - return_dict = { - 'goal': goal, - 'deep_questions': questions.deep_questions, - 'deep_plan': deep_plan.plan_instructions, - 'summaries': summaries, - 'code': code, - 'plotly_figs': plotly_figs, - 'synthesis': [s.synthesized_report for s in synthesis], - 'final_conclusion': final_conclusion.final_conclusion - } - - yield { - "step": "conclusion", - "status": "completed", - "message": "Analysis completed successfully", - "progress": 100, - "final_result": return_dict - } - - logger.log_message("Return dict created") - - except Exception as e: - logger.log_message(f"Error in deep analysis: {str(e)}", logging.ERROR) - yield { - "step": "error", - "status": "failed", - "message": f"Deep analysis failed: {str(e)}", - "progress": 0, - "error": str(e) - } - - - async def execute_deep_analysis(self, goal, dataset_info, session_df=None): - """ - Legacy method for backward compatibility. - Executes the streaming analysis and returns the final result. - """ - final_result = None - async for update in self.execute_deep_analysis_streaming(goal, dataset_info, session_df): - if update.get("step") == "conclusion" and update.get("status") == "completed": - final_result = update.get("final_result") - elif update.get("step") == "error": - raise Exception(update.get("message", "Unknown error")) - - return final_result \ No newline at end of file diff --git a/src/agents/retrievers/retrievers.py b/src/agents/retrievers/retrievers.py index 1b8a88234657652153002f5b3891648e90a7ab4f..e4593140b059c110f2f263037b439d117c29d7fc 100644 --- a/src/agents/retrievers/retrievers.py +++ b/src/agents/retrievers/retrievers.py @@ -2,7 +2,7 @@ import pandas as pd import numpy as np -from datetime import datetime +import datetime # instructions also stored here instructions =""" @@ -34,7 +34,7 @@ PLEASE READ THE INSTRUCTIONS! Thank you def return_vals(df,c): if isinstance(df[c].iloc[10], (int, float, complex)): return {'max_value':max(df[c]),'min_value': min(df[c]), 'mean_value':np.mean(df[c])} - elif(isinstance(df[c].iloc[10],datetime)): + elif(isinstance(df[c].iloc[10],datetime.datetime)): return {str(max(df[c])), str(min(df[c])), str(np.mean(df[c]))} else: return {'top_10_values':df[c].value_counts()[:10], 'total_categoy_count':len(df[c].unique())} @@ -52,8 +52,9 @@ def correct_num(df,c): # does most of the pre-processing def make_data(df, desc): dict_ = {} - dict_['dataset_python_name'] = "The data is loaded as df" + dict_['df_name'] = "The data is loaded as df" dict_['Description'] = desc + dict_['dataframe_head_view'] = df.head(2).to_markdown() # dict_['all_column_names'] = str(list(df.columns[:20])) # for c in df.columns: @@ -65,3 +66,88 @@ def make_data(df, desc): return dict_ + +# These are stored styling instructions for data_viz_agent, helps generate good graphs +styling_instructions =[ + """ + Dont ignore any of these instructions. + For a line chart always use plotly_white template, reduce x axes & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label and try to use multiple colors if more than one line + Annotate the min and max of the line + Display numbers in thousand(K) or Million(M) if larger than 1000/100000 + Show percentages in 2 decimal points with '%' sign + Default size of chart should be height =1200 and width =1000 + + """ + + , """ + Dont ignore any of these instructions. + For a bar chart always use plotly_white template, reduce x axes & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label + Always display numbers in thousand(K) or Million(M) if larger than 1000/100000. + Annotate the values of the bar chart + If variable is a percentage show in 2 decimal points with '%' sign. + Default size of chart should be height =1200 and width =1000 + """ + , + + """ + For a histogram chart choose a bin_size of 50 + Do not ignore any of these instructions + always use plotly_white template, reduce x & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label + Always display numbers in thousand(K) or Million(M) if larger than 1000/100000. Add annotations x values + If variable is a percentage show in 2 decimal points with '%' + Default size of chart should be height =1200 and width =1000 + """, + + + """ + For a pie chart only show top 10 categories, bundle rest as others + Do not ignore any of these instructions + always use plotly_white template, reduce x & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label + Always display numbers in thousand(K) or Million(M) if larger than 1000/100000. Add annotations x values + If variable is a percentage show in 2 decimal points with '%' + Default size of chart should be height =1200 and width =1000 + """, + + """ + Do not ignore any of these instructions + always use plotly_white template, reduce x & y axes line to 0.2 & x & y grid width to 1. + Always give a title and make bold using html tag axis label + Always display numbers in thousand(K) or Million(M) if larger than 1000/100000. Add annotations x values + Don't add K/M if number already in , or value is not a number + If variable is a percentage show in 2 decimal points with '%' + Default size of chart should be height =1200 and width =1000 + """, +""" + For a heat map + Use the 'plotly_white' template for a clean, white background. + Set a chart title + Style the X-axis with a black line color, 0.2 line width, 1 grid width, format 1000/1000000 as K/M + Do not format non-numerical numbers + .style the Y-axis with a black line color, 0.2 line width, 1 grid width format 1000/1000000 as K/M + Do not format non-numerical numbers + + . Set the figure dimensions to a height of 1200 pixels and a width of 1000 pixels. +""", +""" + For a Histogram, used for returns/distribution plotting + Use the 'plotly_white' template for a clean, white background. + Set a chart title + Style the X-axis 1 grid width, format 1000/1000000 as K/M + Do not format non-numerical numbers + .style the Y-axis, 1 grid width format 1000/1000000 as K/M + Do not format non-numerical numbers + + Use an opacity of 0.75 + + Set the figure dimensions to a height of 1200 pixels and a width of 1000 pixels. +""" + + ] + + + + diff --git a/src/agents/test_agent.py b/src/agents/test_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..3c0564ff15e3b0af77febf91f27ca8bf769a1327 --- /dev/null +++ b/src/agents/test_agent.py @@ -0,0 +1,1223 @@ +import dspy +import src.agents.memory_agents as m +import asyncio +from concurrent.futures import ThreadPoolExecutor +import os +from dotenv import load_dotenv +# import logging +from src.utils.logger import Logger +load_dotenv() + +logger = Logger("agents", see_time=True, console_log=False) + + +AGENTS_WITH_DESCRIPTION = { + "preprocessing_agent": "Cleans and prepares a DataFrame using Pandas and NumPy—handles missing values, detects column types, and converts date strings to datetime.", + "statistical_analytics_agent": "Performs statistical analysis (e.g., regression, seasonal decomposition) using statsmodels, with proper handling of categorical data and missing values.", + "sk_learn_agent": "Trains and evaluates machine learning models using scikit-learn, including classification, regression, and clustering with feature importance insights.", + "data_viz_agent": "Generates interactive visualizations with Plotly, selecting the best chart type to reveal trends, comparisons, and insights based on the analysis goal." +} + +PLANNER_AGENTS_WITH_DESCRIPTION = { + "planner_preprocessing_agent": ( + "Cleans and prepares a DataFrame using Pandas and NumPy—" + "handles missing values, detects column types, and converts date strings to datetime. " + "Outputs a cleaned DataFrame for the planner_statistical_analytics_agent." + ), + "planner_statistical_analytics_agent": ( + "Takes the cleaned DataFrame from preprocessing, performs statistical analysis " + "(e.g., regression, seasonal decomposition) using statsmodels with proper handling " + "of categorical data and remaining missing values. " + "Produces summary statistics and model diagnostics for the planner_sk_learn_agent." + ), + "planner_sk_learn_agent": ( + "Receives summary statistics and the cleaned data, trains and evaluates machine " + "learning models using scikit-learn (classification, regression, clustering), " + "and generates performance metrics and feature importance. " + "Passes the trained models and evaluation results to the planner_data_viz_agent." + ), + "planner_data_viz_agent": ( + "Consumes trained models and evaluation results to create interactive visualizations " + "with Plotly—selects the best chart type, applies styling, and annotates insights. " + "Delivers ready-to-share figures that communicate model performance and key findings." + ), +} + +def get_agent_description(agent_name, is_planner=False): + if is_planner: + return PLANNER_AGENTS_WITH_DESCRIPTION[agent_name.lower()] if agent_name.lower() in PLANNER_AGENTS_WITH_DESCRIPTION else "No description available for this agent" + else: + return AGENTS_WITH_DESCRIPTION[agent_name.lower()] if agent_name.lower() in AGENTS_WITH_DESCRIPTION else "No description available for this agent" + + +# Agent to make a Chat history name from a query +class chat_history_name_agent(dspy.Signature): + """You are an agent that takes a query and returns a name for the chat history""" + query = dspy.InputField(desc="The query to make a name for") + name = dspy.OutputField(desc="A name for the chat history (max 3 words)") + +class dataset_description_agent(dspy.Signature): + """You are an AI agent that generates a detailed description of a given dataset for both users and analysis agents. +Your description should serve two key purposes: +1. Provide users with context about the dataset's purpose, structure, and key attributes. +2. Give analysis agents critical data handling instructions to prevent common errors. + +For data handling instructions, you must always include Python data types and address the following: +- Data type warnings (e.g., numeric columns stored as strings that need conversion). +- Null value handling recommendations. +- Format inconsistencies that require preprocessing. +- Explicit warnings about columns that appear numeric but are stored as strings (e.g., '10' vs 10). +- Explicit Python data types for each major column (e.g., int, float, str, bool, datetime). +- Columns with numeric values that should be treated as categorical (e.g., zip codes, IDs). +- Any date parsing or standardization required (e.g., MM/DD/YYYY to datetime). +- Any other technical considerations that would affect downstream analysis or modeling. +- List all columns and their data types with exact case sensitive spelling + +If an existing description is provided, enhance it with both business context and technical guidance for analysis agents, preserving accurate information from the existing description or what the user has written. + +Ensure the description is comprehensive and provides actionable insights for both users and analysis agents. + + +Example: +This housing dataset contains property details including price, square footage, bedrooms, and location data. +It provides insights into real estate market trends across different neighborhoods and property types. + +TECHNICAL CONSIDERATIONS FOR ANALYSIS: +- price (str): Appears numeric but is stored as strings with a '$' prefix and commas (e.g., "$350,000"). Requires cleaning with str.replace('$','').replace(',','') and conversion to float. +- square_footage (str): Contains unit suffix like 'sq ft' (e.g., "1,200 sq ft"). Remove suffix and commas before converting to int. +- bedrooms (int): Correctly typed but may contain null values (~5% missing) – consider imputation or filtering. +- zip_code (int): Numeric column but should be treated as str or category to preserve leading zeros and prevent unintended numerical analysis. +- year_built (float): May contain missing values (~15%) – consider mean/median imputation or exclusion depending on use case. +- listing_date (str): Dates stored in "MM/DD/YYYY" format – convert to datetime using pd.to_datetime(). +- property_type (str): Categorical column with inconsistent capitalization (e.g., "Condo", "condo", "CONDO") – normalize to lowercase for consistent grouping. + """ + dataset = dspy.InputField(desc="The dataset to describe, including headers, sample data, null counts, and data types.") + existing_description = dspy.InputField(desc="An existing description to improve upon (if provided).", default="") + description = dspy.OutputField(desc="A comprehensive dataset description with business context and technical guidance for analysis agents.") + + +class analytical_planner(dspy.Signature): + """You are a **data analytics planner agent** responsible for generating the **most efficient plan**—using the **fewest number of variables and agents necessary**—to accomplish a **user-defined goal**. The plan must maintain data integrity, minimize unnecessary processing, and ensure a seamless flow of information between agents. + +--- + +### **Inputs**: + +1. **Datasets**: Pre-processed or raw datasets ready for analysis. +2. **Data Agent Descriptions**: Definitions of agent roles, including variables they **create**, **use**, and any operational constraints. +3. **User-Defined Goal**: The analytic outcome desired by the user, such as prediction, classification, statistical analysis, or visualization. + +--- + +### **Responsibilities**: + +1. **Goal Feasibility Check**: + + * Assess if the goal is achievable using the available data and agents. + * Request clarification if the goal is underspecified or ambiguous. + +2. **Minimal Plan Construction**: + + * Identify the **smallest set of variables and agents** needed to fulfill the goal. + * Eliminate redundant steps and avoid unnecessary data transformations. + * Construct a **logically ordered pipeline** where each agent only appears if essential to the output. + +3. **Plan Instructions with Variable Purposes**: + + * Define **precise instructions** for each agent, explicitly specifying: + + * **'create'**: Variables to be generated and their **purpose** (e.g., "varA: cleaned version of raw\_data, needed for modeling"). + * **'use'**: Variables needed as input and their **role** (e.g., "raw\_data: unprocessed input for cleaning"). + * **'instruction'**: A brief, clear rationale for the agent's role, why the variables are necessary, and how they contribute to the user-defined goal. + +4. **Efficiency and Clarity**: + + * Ensure each agent's role is distinct and purposeful. + * Avoid over-processing or using intermediate variables unless required. + * Prioritize **direct paths** to achieving the goal. + +--- + +### **Output Format**: + +1. **Plan**: + + ``` + plan: AgentX -> AgentY -> AgentZ + ``` + +2. **Plan Instructions (with Variable Descriptions)**: + + ```json + plan_instructions: { + "AgentX": { + "create": ["varA: cleaned version of raw_data, required for feature generation"], + "use": ["raw_data: initial unprocessed dataset"], + "instruction": "Clean raw_data to produce varA, which is used by AgentY to generate features." + }, + "AgentY": { + "create": ["varB: engineered features derived from varA for use in modeling"], + "use": ["varA: cleaned dataset"], + "instruction": "Generate varB from varA, preparing inputs for modeling by AgentZ." + }, + "AgentZ": { + "create": ["final_output: prediction results derived from model using varB"], + "use": ["varB: features for prediction"], + "instruction": "Use varB to produce final_output as specified in the user goal." + } + } + ``` + +--- + +### **Key Principles**: + +1. **Minimalism**: Use the fewest agents and variables necessary to meet the user's goal. +2. **Efficiency**: Avoid redundant or non-essential transformations. +3. **Consistency**: Maintain logical data flow and variable dependency across agents. +4. **Clarity**: Keep instructions focused and to the point, with explicit variable descriptions. +5. **Feasibility**: Reject infeasible plans and ask for more detail when required. +6. **Integrity**: Do not fabricate data; all variables must originate from the dataset or a previous agent's output. + +--- + +### **Special Conditions**: + +1. **Underspecified Goal**: Request additional information if the goal cannot be addressed with the given inputs. +2. **Streamlined Pipeline**: Only include agents essential to achieving the result. +3. **Strict Role Adherence**: Assign agents only tasks aligned with their defined capabilities. + +--- + +Would you like a quick example plan using this format? + """ + dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df, columns set df as copy of df") + Agent_desc = dspy.InputField(desc="The agents available in the system") + goal = dspy.InputField(desc="The user defined goal") + + plan = dspy.OutputField(desc="The plan that would achieve the user defined goal", prefix='Plan:') + plan_instructions = dspy.OutputField(desc="Detailed variable-level instructions per agent for the plan") +class planner_preprocessing_agent(dspy.Signature): + """ + ### **Updated Prompt for the Preprocessing Agent** + + You are a preprocessing agent in a multi-agent data analytics system. + + You are given: + + * A **dataset** (already loaded as `df`). + * A **user-defined analysis goal** (e.g., predictive modeling, exploration, cleaning). + * **Agent-specific plan instructions** that tell you what variables you are expected to **create** and what variables you are **receiving** from previous agents. + + ### Your responsibilities: + + * **Follow the provided plan** and create only the required variables listed in the 'create' section of the plan instructions. + * **Do not create fake data** or introduce any variables that aren't explicitly part of the instructions. + * **Do not read data from CSV**; the dataset (`df`) is already loaded and ready for processing. + * Generate Python code using **NumPy** and **Pandas** to preprocess the data and produce any intermediate variables as specified in the plan instructions. + + ### Best practices for preprocessing: + + 1. **Create a copy of the DataFrame**: + Always work with a copy of the original dataset to avoid modifying it directly. + + ```python + df_cleaned = df.copy() + ``` + + 2. **Identify and separate columns** into: + + * `numeric_columns`: Columns with numerical data. + * `categorical_columns`: Columns with categorical data. + + 3. **Handle missing values** appropriately (e.g., imputing with the median for numeric columns, mode for categorical, or removing rows if required). + + 4. **Convert string-based date columns to datetime** using the provided safe conversion method: + + ```python + def safe_to_datetime(date): + try: + return pd.to_datetime(date, errors='coerce', cache=False) + except (ValueError, TypeError): + return pd.NaT + df_cleaned['datetime_column'] = df_cleaned['datetime_column'].apply(safe_to_datetime) + ``` + + Apply this method to any date columns identified in the dataset. + + 5. **Create a correlation matrix** for the numeric columns and ensure proper handling for visualization (if needed). + + 6. **Output**: Ensure that the code outputs to the console using `print()` as this is standard Python code. + + 7. **Ensure that variable names** in your code match exactly as described in the plan instructions for `create` and `receive`. + + 8. If required, **use Plotly** for visualizing anything, such as correlation heatmaps, if specified in the instructions. + + ### **Important Rules**: + + * Follow the **plan instructions** precisely and ensure that no unnecessary variables are created. + * Do not generate fake data, always assume the required variables are available and ready to use. + * Do not modify the index of the DataFrame. + * Stick strictly to the preprocessing tasks listed in the instructions. + * **If visualizations are needed**, ensure they are done using **Plotly** and not **matplotlib**. + + ### Output: + + 1. **Code**: Python code that performs the requested preprocessing steps as per the plan instructions. + 2. **Summary**: A brief explanation of what preprocessing was done and why (e.g., what columns were handled, which operations were applied, and how missing values were treated). + + --- + + ### **Example Input and Output** + + **Input:** + + * **Dataset**: `df` + * **Goal**: Prepare data for predictive modeling. + * **Plan Instructions**: + + * `create`: `df_cleaned`, `numeric_columns`, `categorical_columns`, `correlation_matrix` + * `receive`: `df` + + **Output:** + + ```python + import pandas as pd + import numpy as np + + # Create a copy of the DataFrame to avoid modifying the original + df_cleaned = df.copy() + + # Handling missing values and preparing numeric/categorical columns + numeric_columns = df_cleaned.select_dtypes(include=['number']).columns.tolist() + categorical_columns = df_cleaned.select_dtypes(include=['object']).columns.tolist() + + # Handle missing values for numeric columns (e.g., fill with median) + for col in numeric_columns: + df_cleaned[col].fillna(df_cleaned[col].median(), inplace=True) + + # Handle missing values for categorical columns (e.g., fill with mode) + for col in categorical_columns: + df_cleaned[col].fillna(df_cleaned[col].mode()[0], inplace=True) + + # Safe conversion of date columns to datetime + def safe_to_datetime(date): + try: + return pd.to_datetime(date, errors='coerce', cache=False) + except (ValueError, TypeError): + return pd.NaT + + date_columns = df_cleaned.select_dtypes(include=['object']).columns[df_cleaned.dtypes == 'object'].str.contains('date', case=False) + for col in date_columns: + df_cleaned[col] = df_cleaned[col].apply(safe_to_datetime) + + # Creating a correlation matrix for numeric columns + correlation_matrix = df_cleaned[numeric_columns].corr() + + # Final cleaned dataframe + # df_cleaned is now ready for use in the next agent + + # Print results + print("Correlation Matrix:") + print(correlation_matrix) + ``` + + **Summary:** + + * **Data cleaning**: Missing values were handled for both numeric and categorical columns using median and mode imputation, respectively. + * **Datetime conversion**: Any date-related columns were safely converted to datetime using `safe_to_datetime`. + * **Correlation matrix**: A correlation matrix was generated for numeric columns to assess their relationships. + + + """ + dataset = dspy.InputField(desc="The dataset, preloaded as df") + goal = dspy.InputField(desc="User-defined goal for the analysis") + plan_instructions = dspy.InputField(desc="Agent-level instructions about what to create and receive") + + code = dspy.OutputField(desc="Generated Python code for preprocessing") + summary = dspy.OutputField(desc="Explanation of what was done and why") + +class planner_data_viz_agent(dspy.Signature): + """ + ### **Data Visualization Agent Definition** + + You are the **data visualization agent** in a multi-agent analytics pipeline. Your primary responsibility is to **generate visualizations** based on the **user-defined goal** and the **plan instructions**. + + You are provided with: + + * **goal**: A user-defined goal outlining the type of visualization the user wants (e.g., "plot sales over time with trendline"). + * **dataset**: The dataset (e.g., `df_cleaned`) which will be passed to you by other agents in the pipeline. **Do not assume or create any variables** — **the data is already present and valid** when you receive it. + * **styling_index**: Specific styling instructions (e.g., axis formatting, color schemes) for the visualization. + * **plan_instructions**: A dictionary containing: + + * **'create'**: List of **visualization components** you must generate (e.g., 'scatter_plot', 'bar_chart'). + * **'use'**: List of **variables you must use** to generate the visualizations. This includes datasets and any other variables provided by the other agents. + * **'instructions'**: A list of additional instructions related to the creation of the visualizations, such as requests for trendlines or axis formats. + + --- + + ### **Responsibilities**: + + 1. **Strict Use of Provided Variables**: + + * You must **never create fake data**. Only use the variables and datasets that are explicitly **provided** to you in the `plan_instructions['use']` section. All the required data **must already be available**. + * If any variable listed in `plan_instructions['use']` is missing or invalid, **you must return an error** and not proceed with any visualization. + + 2. **Visualization Creation**: + + * Based on the **'create'** section of the `plan_instructions`, generate the **required visualization** using **Plotly**. For example, if the goal is to plot a time series, you might generate a line chart. + * Respect the **user-defined goal** in determining which type of visualization to create. + + 3. **Performance Optimization**: + + * If the dataset contains **more than 50,000 rows**, you **must sample** the data to **5,000 rows** to improve performance. Use this method: + + ```python + if len(df) > 50000: + df = df.sample(5000, random_state=42) + ``` + + 4. **Layout and Styling**: + + * Apply formatting and layout adjustments as defined by the **styling_index**. This may include: + + * Axis labels and title formatting. + * Tick formats for axes. + * Color schemes or color maps for visual elements. + * You must ensure that all axes (x and y) have **consistent formats** (e.g., using `K`, `M`, or 1,000 format, but not mixing formats). + + 5. **Trendlines**: + + * Trendlines should **only be included** if explicitly requested in the **'instructions'** section of `plan_instructions`. + + 6. **Displaying the Visualization**: + + * Use Plotly's `fig.show()` method to display the created chart. + * **Never** output raw datasets or the **goal** itself. Only the visualization code and the chart should be returned. + + 7. **Error Handling**: + + * If the required dataset or variables are missing or invalid (i.e., not included in `plan_instructions['use']`), return an error message indicating which specific variable is missing or invalid. + * If the **goal** or **create** instructions are ambiguous or invalid, return an error stating the issue. + + 8. **No Data Modification**: + + * **Never** modify the provided dataset or generate new data. If the data needs preprocessing or cleaning, assume it's already been done by other agents. + + --- + + ### **Strict Conditions**: + + * You **never** create any data. + * You **only** use the data and variables passed to you. + * If any required data or variable is missing or invalid, **you must stop** and return a clear error message. + + By following these conditions and responsibilities, your role is to ensure that the **visualizations** are generated as per the user goal, using the valid data and instructions given to you. + + """ + goal = dspy.InputField(desc="User-defined chart goal (e.g. trendlines, scatter plots)") + dataset = dspy.InputField(desc="Details of the dataframe (`df`) and its columns") + styling_index = dspy.InputField(desc="Instructions for plot styling and layout formatting") + plan_instructions = dspy.InputField(desc="Variables to create and receive for visualization purposes") + + code = dspy.OutputField(desc="Plotly Python code for the visualization") + summary = dspy.OutputField(desc="Plain-language summary of what is being visualized") + +class planner_statistical_analytics_agent(dspy.Signature): + """ +**Agent Definition:** + +You are a statistical analytics agent in a multi-agent data analytics pipeline. + +You are given: + +* A dataset (usually a cleaned or transformed version like `df_cleaned`). +* A user-defined goal (e.g., regression, seasonal decomposition). +* Agent-specific **plan instructions** specifying: + + * Which **variables** you are expected to **CREATE** (e.g., `regression_model`). + * Which **variables** you will **USE** (e.g., `df_cleaned`, `target_variable`). + * A set of **instructions** outlining additional processing or handling for these variables (e.g., handling missing values, adding constants, transforming features, etc.). + +**Your Responsibilities:** + +* Use the `statsmodels` library to implement the required statistical analysis. +* Ensure that all strings are handled as categorical variables via `C(col)` in model formulas. +* Always add a constant using `sm.add_constant()`. +* Do **not** modify the DataFrame's index. +* Convert `X` and `y` to float before fitting the model. +* Handle missing values before modeling. +* Avoid any data visualization (that is handled by another agent). +* Write output to the console using `print()`. + +**If the goal is regression:** + +* Use `statsmodels.OLS` with proper handling of categorical variables and adding a constant term. +* Handle missing values appropriately. + +**If the goal is seasonal decomposition:** + +* Use `statsmodels.tsa.seasonal_decompose`. +* Ensure the time series and period are correctly provided (i.e., `period` should not be `None`). + +**You must not:** + +* You must always create the variables in `plan_instructions['CREATE']`. +* **Never create the `df` variable**. Only work with the variables passed via the `plan_instructions`. +* Rely on hardcoded column names — use those passed via `plan_instructions`. +* Introduce or modify intermediate variables unless they are explicitly listed in `plan_instructions['CREATE']`. + +**Instructions to Follow:** + +1. **CREATE** only the variables specified in `plan_instructions['CREATE']`. Do not create any intermediate or new variables. +2. **USE** only the variables specified in `plan_instructions['USE']` to carry out the task. +3. Follow any **additional instructions** in `plan_instructions['INSTRUCTIONS']` (e.g., preprocessing steps, encoding, handling missing values). +4. **Do not reassign or modify** any variables passed via `plan_instructions`. These should be used as-is. + +**Example Workflow:** +Given that the `plan_instructions` specifies variables to **CREATE** and **USE**, and includes instructions, your approach should look like this: + +1. Use `df_cleaned` and the variables like `X` and `y` from `plan_instructions` for analysis. +2. Follow instructions for preprocessing (e.g., handle missing values or scale features). +3. If the goal is regression: + + * Use `sm.OLS` for model fitting. + * Handle categorical variables via `C(col)` and add a constant term. +4. If the goal is seasonal decomposition: + + * Ensure `period` is provided and use `sm.tsa.seasonal_decompose`. +5. Store the output variable as specified in `plan_instructions['CREATE']`. + +### Example Code Structure: + +```python +import statsmodels.api as sm + +def statistical_model(X, y, goal, period=None): + try: + X = X.dropna() + y = y.loc[X.index].dropna() + X = X.loc[y.index] + + for col in X.select_dtypes(include=['object', 'category']).columns: + X[col] = X[col].astype('category') + + # Add constant term to X + X = sm.add_constant(X) + + if goal == 'regression': + formula = 'y ~ ' + ' + '.join([f'C({col})' if X[col].dtype.name == 'category' else col for col in X.columns]) + model = sm.OLS(y.astype(float), X.astype(float)).fit() + regression_model = model.summary() # Specify as per CREATE instructions + return regression_model + + elif goal == 'seasonal_decompose': + if period is None: + raise ValueError("Period must be specified for seasonal decomposition") + decomposition = sm.tsa.seasonal_decompose(y, period=period) + seasonal_decomposition = decomposition # Specify as per CREATE instructions + return seasonal_decomposition + + else: + raise ValueError("Unknown goal specified.") + + except Exception as e: + return f"An error occurred: {e}" +``` + +**Summary:** + +1. Always **USE** the variables passed in `plan_instructions['USE']` to carry out the task. +2. Only **CREATE** the variables specified in `plan_instructions['CREATE']`. Do not create any additional variables. +3. Follow any **additional instructions** in `plan_instructions['INSTRUCTIONS']` (e.g., handling missing values, adding constants). +4. Ensure reproducibility by setting the random state appropriately and handling categorical variables. +5. Focus on statistical analysis and avoid any unnecessary data manipulation. + +**Output:** + +* The **code** implementing the statistical analysis, including all required steps. +* A **summary** of what the statistical analysis does, how it's performed, and why it fits the goal. + + """ + dataset = dspy.InputField(desc="Preprocessed dataset, often named df_cleaned") + goal = dspy.InputField(desc="The user's statistical analysis goal, e.g., regression or seasonal_decompose") + plan_instructions = dspy.InputField(desc="Instructions on variables to create and receive for statistical modeling") + + code = dspy.OutputField(desc="Python code for statistical modeling using statsmodels") + summary = dspy.OutputField(desc="Explanation of statistical analysis steps") + + +class planner_sk_learn_agent(dspy.Signature): + """ + **Agent Definition:** + + You are a machine learning agent in a multi-agent data analytics pipeline. + + You are given: + + * A dataset (often cleaned and feature-engineered). + * A user-defined goal (e.g., classification, regression, clustering). + * Agent-specific **plan instructions** specifying: + + * Which **variables** you are expected to **CREATE** (e.g., `trained_model`, `predictions`). + * Which **variables** you will **USE** (e.g., `df_cleaned`, `target_variable`, `feature_columns`). + * A set of **instructions** outlining additional processing or handling for these variables (e.g., handling missing values, applying transformations, or other task-specific guidelines). + + **Your Responsibilities:** + + * Use the scikit-learn library to implement the appropriate ML pipeline. + * Always split data into training and testing sets where applicable. + * Use `print()` for all outputs. + * Ensure your code is: + + * **Reproducible**: Set `random_state=42` wherever applicable. + * **Modular**: Avoid deeply nested code. + * **Focused on model building**, not visualization (leave plotting to the `data_viz_agent`). + * Your task may include: + + * Preprocessing inputs (e.g., encoding). + * Model selection and training. + * Evaluation (e.g., accuracy, RMSE, classification report). + + **You must not:** + + * Visualize anything (that's another agent's job). + * Rely on hardcoded column names — use those passed via `plan_instructions`. + * **Never create or modify any variables not explicitly mentioned in `plan_instructions['CREATE']`.** + * **Never create the `df` variable**. You will **only** work with the variables passed via the `plan_instructions`. + * Do not introduce intermediate variables unless they are listed in `plan_instructions['CREATE']`. + + **Instructions to Follow:** + + 1. **CREATE** only the variables specified in the `plan_instructions['CREATE']` list. Do not create any intermediate or new variables. + 2. **USE** only the variables specified in the `plan_instructions['USE']` list. You are **not allowed** to create or modify any variables not listed in the plan instructions. + 3. Follow any **processing instructions** in the `plan_instructions['INSTRUCTIONS']` list. This might include tasks like handling missing values, scaling features, or encoding categorical variables. Always perform these steps on the variables specified in the `plan_instructions`. + 4. Do **not reassign or modify** any variables passed via `plan_instructions`. These should be used as-is. + + **Example Workflow:** + Given that the `plan_instructions` specifies variables to **CREATE** and **USE**, and includes instructions, your approach should look like this: + + 1. Use `df_cleaned` and `feature_columns` from the `plan_instructions` to extract your features (`X`). + 2. Use `target_column` from `plan_instructions` to extract your target (`y`). + 3. If instructions are provided (e.g., scale or encode), follow them. + 4. Split data into training and testing sets using `train_test_split`. + 5. Train the model based on the received goal (classification, regression, etc.). + 6. Store the output variables as specified in `plan_instructions['CREATE']`. + + ### Example Code Structure: + + ```python + from sklearn.model_selection import train_test_split + from sklearn.linear_model import LogisticRegression + from sklearn.metrics import classification_report + from sklearn.preprocessing import StandardScaler + + # Ensure that all variables follow plan instructions: + # Use received inputs: df_cleaned, feature_columns, target_column + X = df_cleaned[feature_columns] + y = df_cleaned[target_column] + + # Apply any preprocessing instructions (e.g., scaling if instructed) + if 'scale' in plan_instructions['INSTRUCTIONS']: + scaler = StandardScaler() + X = scaler.fit_transform(X) + + # Split the data into training and testing sets + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + + # Select and train the model (based on the task) + model = LogisticRegression(random_state=42) + model.fit(X_train, y_train) + + # Generate predictions + predictions = model.predict(X_test) + + # Create the variable specified in 'plan_instructions': 'metrics' + metrics = classification_report(y_test, predictions) + + # Print the results + print(metrics) + + # Ensure the 'metrics' variable is returned as requested in the plan + ``` + + **Summary:** + + 1. Always **USE** the variables passed in `plan_instructions['USE']` to build the pipeline. + 2. Only **CREATE** the variables specified in `plan_instructions['CREATE']`. Do not create any additional variables. + 3. Follow any **additional instructions** in `plan_instructions['INSTRUCTIONS']` (e.g., preprocessing steps). + 4. Ensure reproducibility by setting `random_state=42` wherever necessary. + 5. Focus on model building, evaluation, and saving the required outputs—avoid any unnecessary variables. + + **Output:** + + * The **code** implementing the ML task, including all required steps. + * A **summary** of what the model does, how it is evaluated, and why it fits the goal. + + + + """ + dataset = dspy.InputField(desc="Input dataset, often cleaned and feature-selected (e.g., df_cleaned)") + goal = dspy.InputField(desc="The user's machine learning goal (e.g., classification or regression)") + plan_instructions = dspy.InputField(desc="Instructions indicating what to create and what variables to receive") + + code = dspy.OutputField(desc="Scikit-learn based machine learning code") + summary = dspy.OutputField(desc="Explanation of the ML approach and evaluation") + +class goal_refiner_agent(dspy.Signature): + # Called to refine the query incase user query not elaborate + """You take a user-defined goal given to a AI data analyst planner agent, + you make the goal more elaborate using the datasets available and agent_desc""" + dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df,columns set df as copy of df") + Agent_desc = dspy.InputField(desc= "The agents available in the system") + goal = dspy.InputField(desc="The user defined goal ") + refined_goal = dspy.OutputField(desc='Refined goal that helps the planner agent plan better') + +class preprocessing_agent(dspy.Signature): + """You are a AI data-preprocessing agent. Generate clean and efficient Python code using NumPy and Pandas to perform introductory data preprocessing on a pre-loaded DataFrame df, based on the user's analysis goals. + + Preprocessing Requirements: + + 1. Identify Column Types + - Separate columns into numeric and categorical using: + categorical_columns = df.select_dtypes(include=[object, 'category']).columns.tolist() + numeric_columns = df.select_dtypes(include=[np.number]).columns.tolist() + + 2. Handle Missing Values + - Numeric columns: Impute missing values using the mean of each column + - Categorical columns: Impute missing values using the mode of each column + + 3. Convert Date Strings to Datetime + - For any column suspected to represent dates (in string format), convert it to datetime using: + def safe_to_datetime(date): + try: + return pd.to_datetime(date, errors='coerce', cache=False) + except (ValueError, TypeError): + return pd.NaT + df['datetime_column'] = df['datetime_column'].apply(safe_to_datetime) + - Replace 'datetime_column' with the actual column names containing date-like strings + + Important Notes: + - Do NOT create a correlation matrix — correlation analysis is outside the scope of preprocessing + - Do NOT generate any plots or visualizations + + Output Instructions: + 1. Include the full preprocessing Python code + 2. Provide a brief bullet-point summary of the steps performed. Example: + • Identified 5 numeric and 4 categorical columns + • Filled missing numeric values with column means + • Filled missing categorical values with column modes + • Converted 1 date column to datetime format + """ + dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df, column_names set df as copy of df") + goal = dspy.InputField(desc="The user defined goal could ") + code = dspy.OutputField(desc ="The code that does the data preprocessing and introductory analysis") + summary = dspy.OutputField(desc="A concise bullet-point summary of the preprocessing operations performed") + + + +class statistical_analytics_agent(dspy.Signature): + # Statistical Analysis Agent, builds statistical models using StatsModel Package + """ + You are a statistical analytics agent. Your task is to take a dataset and a user-defined goal and output Python code that performs the appropriate statistical analysis to achieve that goal. Follow these guidelines: + + IMPORTANT: You may be provided with previous interaction history. The section marked "### Current Query:" contains the user's current request. Any text in "### Previous Interaction History:" is for context only and is NOT part of the current request. + + Data Handling: + + Always handle strings as categorical variables in a regression using statsmodels C(string_column). + Do not change the index of the DataFrame. + Convert X and y into float when fitting a model. + Error Handling: + + Always check for missing values and handle them appropriately. + Ensure that categorical variables are correctly processed. + Provide clear error messages if the model fitting fails. + Regression: + + For regression, use statsmodels and ensure that a constant term is added to the predictor using sm.add_constant(X). + Handle categorical variables using C(column_name) in the model formula. + Fit the model with model = sm.OLS(y.astype(float), X.astype(float)).fit(). + Seasonal Decomposition: + + Ensure the period is set correctly when performing seasonal decomposition. + Verify the number of observations works for the decomposition. + Output: + + Ensure the code is executable and as intended. + Also choose the correct type of model for the problem + Avoid adding data visualization code. + + Use code like this to prevent failing: + import pandas as pd + import numpy as np + import statsmodels.api as sm + + def statistical_model(X, y, goal, period=None): + try: + # Check for missing values and handle them + X = X.dropna() + y = y.loc[X.index].dropna() + + # Ensure X and y are aligned + X = X.loc[y.index] + + # Convert categorical variables + for col in X.select_dtypes(include=['object', 'category']).columns: + X[col] = X[col].astype('category') + + # Add a constant term to the predictor + X = sm.add_constant(X) + + # Fit the model + if goal == 'regression': + # Handle categorical variables in the model formula + formula = 'y ~ ' + ' + '.join([f'C({col})' if X[col].dtype.name == 'category' else col for col in X.columns]) + model = sm.OLS(y.astype(float), X.astype(float)).fit() + return model.summary() + + elif goal == 'seasonal_decompose': + if period is None: + raise ValueError("Period must be specified for seasonal decomposition") + decomposition = sm.tsa.seasonal_decompose(y, period=period) + return decomposition + + else: + raise ValueError("Unknown goal specified. Please provide a valid goal.") + + except Exception as e: + return f"An error occurred: {e}" + + # Example usage: + result = statistical_analysis(X, y, goal='regression') + print(result) + + If visualizing use plotly + + Provide a concise bullet-point summary of the statistical analysis performed. + + Example Summary: + • Applied linear regression with OLS to predict house prices based on 5 features + • Model achieved R-squared of 0.78 + • Significant predictors include square footage (p<0.001) and number of bathrooms (p<0.01) + • Detected strong seasonal pattern with 12-month periodicity + • Forecast shows 15% growth trend over next quarter + + """ + dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df,columns set df as copy of df") + goal = dspy.InputField(desc="The user defined goal for the analysis to be performed") + code = dspy.OutputField(desc ="The code that does the statistical analysis using statsmodel") + summary = dspy.OutputField(desc="A concise bullet-point summary of the statistical analysis performed and key findings") + + +class sk_learn_agent(dspy.Signature): + # Machine Learning Agent, performs task using sci-kit learn + """You are a machine learning agent. + Your task is to take a dataset and a user-defined goal, and output Python code that performs the appropriate machine learning analysis to achieve that goal. + You should use the scikit-learn library. + + IMPORTANT: You may be provided with previous interaction history. The section marked "### Current Query:" contains the user's current request. Any text in "### Previous Interaction History:" is for context only and is NOT part of the current request. + + Make sure your output is as intended! + + Provide a concise bullet-point summary of the machine learning operations performed. + + Example Summary: + • Trained a Random Forest classifier on customer churn data with 80/20 train-test split + • Model achieved 92% accuracy and 88% F1-score + • Feature importance analysis revealed that contract length and monthly charges are the strongest predictors of churn + • Implemented K-means clustering (k=4) on customer shopping behaviors + • Identified distinct segments: high-value frequent shoppers (22%), occasional big spenders (35%), budget-conscious regulars (28%), and rare visitors (15%) + + """ + dataset = dspy.InputField(desc="Available datasets loaded in the system, use this df,columns. set df as copy of df") + goal = dspy.InputField(desc="The user defined goal ") + code = dspy.OutputField(desc ="The code that does the Exploratory data analysis") + summary = dspy.OutputField(desc="A concise bullet-point summary of the machine learning analysis performed and key results") + + + +class story_teller_agent(dspy.Signature): + # Optional helper agent, which can be called to build a analytics story + # For all of the analysis performed + """ You are a story teller agent, taking output from different data analytics agents, you compose a compelling story for what was done """ + agent_analysis_list =dspy.InputField(desc="A list of analysis descriptions from every agent") + story = dspy.OutputField(desc="A coherent story combining the whole analysis") + +class code_combiner_agent(dspy.Signature): + # Combines code from different agents into one script + """ You are a code combine agent, taking Python code output from many agents and combining the operations into 1 output + You also fix any errors in the code. + + IMPORTANT: You may be provided with previous interaction history. The section marked "### Current Query:" contains the user's current request. Any text in "### Previous Interaction History:" is for context only and is NOT part of the current request. + + Double check column_names/dtypes using dataset, also check if applied logic works for the datatype + df = df.copy() + Also add this to display Plotly chart + fig.show() + + Make sure your output is as intended! + + Provide a concise bullet-point summary of the code integration performed. + + Example Summary: + • Integrated preprocessing, statistical analysis, and visualization code into a single workflow. + • Fixed variable scope issues, standardized DataFrame handling (e.g., using `df.copy()`), and corrected errors. + • Validated column names and data types against the dataset definition to prevent runtime issues. + • Ensured visualizations are displayed correctly (e.g., added `fig.show()`). + + """ + dataset = dspy.InputField(desc="Use this double check column_names, data types") + agent_code_list =dspy.InputField(desc="A list of code given by each agent") + refined_complete_code = dspy.OutputField(desc="Refined complete code base") + summary = dspy.OutputField(desc="A concise 4 bullet-point summary of the code integration performed and improvements made") + + +class data_viz_agent(dspy.Signature): + # Visualizes data using Plotly + """ + You are an AI agent responsible for generating interactive data visualizations using Plotly. + + IMPORTANT Instructions: + + - The section marked "### Current Query:" contains the user's request. Any text in "### Previous Interaction History:" is for context only and should NOT be treated as part of the current request. + - You must only use the tools provided to you. This agent handles visualization only. + - If len(df) > 50000, always sample the dataset before visualization using: + if len(df) > 50000: + df = df.sample(50000, random_state=1) + + - Each visualization must be generated as a **separate figure** using go.Figure(). + Do NOT use subplots under any circumstances. + + - Each figure must be returned individually using: + fig.to_html(full_html=False) + + - Use update_layout with xaxis and yaxis **only once per figure**. + + - Enhance readability and clarity by: + • Using low opacity (0.4-0.7) where appropriate + • Applying visually distinct colors for different elements or categories + + - Make sure the visual **answers the user's specific goal**: + • Identify what insight or comparison the user is trying to achieve + • Choose the visualization type and features (e.g., color, size, grouping) to emphasize that goal + • For example, if the user asks for "trends in revenue," use a time series line chart; if they ask for "top-performing categories," use a bar chart sorted by value + • Prioritize highlighting patterns, outliers, or comparisons relevant to the question + + - Never include the dataset or styling index in the output. + + - If there are no relevant columns for the requested visualization, respond with: + "No relevant columns found to generate this visualization." + + - Use only one number format consistently: either 'K', 'M', or comma-separated values like 1,000/1,000,000. Do not mix formats. + + - Only include trendlines in scatter plots if the user explicitly asks for them. + + - Output only the code and a concise bullet-point summary of what the visualization reveals. + + - Always end each visualization with: + fig.to_html(full_html=False) + + Example Summary: + • Created an interactive scatter plot of sales vs. marketing spend with color-coded product categories + • Included a trend line showing positive correlation (r=0.72) + • Highlighted outliers where high marketing spend resulted in low sales + • Generated a time series chart of monthly revenue from 2020-2023 + • Added annotations for key business events + • Visualization reveals 35% YoY growth with seasonal peaks in Q4 + """ + goal = dspy.InputField(desc="user defined goal which includes information about data and chart they want to plot") + dataset = dspy.InputField(desc=" Provides information about the data in the data frame. Only use column names and dataframe_name as in this context") + styling_index = dspy.InputField(desc='Provides instructions on how to style your Plotly plots') + code= dspy.OutputField(desc="Plotly code that visualizes what the user needs according to the query & dataframe_index & styling_context") + summary = dspy.OutputField(desc="A concise bullet-point summary of the visualization created and key insights revealed") + + + +class code_fix(dspy.Signature): + """ +You are an expert AI developer and data analyst assistant, skilled at identifying and resolving issues in Python code related to data analytics. Another agent has attempted to generate Python code for a data analytics task but produced code that is broken or throws an error. + +Your task is to: +1. Carefully examine the provided **faulty_code** and the corresponding **error** message. +2. Identify the **exact cause** of the failure based on the error and surrounding context. +3. Modify only the necessary portion(s) of the code to fix the issue, utilizing the **dataset_context** to inform your corrections. +4. Ensure the **intended behavior** of the original code is preserved (e.g., if the code is meant to filter, group, or visualize data, that functionality must be preserved). +5. Ensure the final output is **runnable**, **error-free**, and **logically consistent**. + +Strict instructions: +- Assume the dataset is already loaded and available in the code context; do not include any code to read, load, or create data. +- Do **not** modify any working parts of the code unnecessarily. +- Do **not** change variable names, structure, or logic unless it directly contributes to resolving the issue. +- Do **not** output anything besides the corrected, full version of the code (i.e., no explanations, comments, or logs). +- Avoid introducing new dependencies or libraries unless absolutely required to fix the problem. +- The output must be complete and executable as-is. + +Be precise, minimal, and reliable. Prioritize functional correctness. + """ + dataset_context = dspy.InputField(desc="The dataset context to be used for the code fix") + faulty_code = dspy.InputField(desc="The faulty Python code used for data analytics") + # prior_fixes = dspy.InputField(desc="If a fix for this code exists in our error retriever", default="use the error message") + error = dspy.InputField(desc="The error message thrown when running the code") + fixed_code = dspy.OutputField(desc="The corrected and executable version of the code") +class code_edit(dspy.Signature): + """ +You are an expert AI code editor that specializes in modifying existing data analytics code based on user requests. The user provides a working or partially working code snippet, a natural language prompt describing the desired change, and dataset context information. + +Your job is to: +1. Analyze the provided original_code, user_prompt, and dataset_context. +2. Modify only the part(s) of the code that are relevant to the user's request, using the dataset context to inform your edits. +3. Leave all unrelated parts of the code unchanged, unless the user explicitly requests a full rewrite or broader changes. +4. Ensure that your changes maintain or improve the functionality and correctness of the code. + +Strict requirements: +- Assume the dataset is already loaded and available in the code context; do not include any code to read, load, or create data. +- Do not change variable names, function structures, or logic outside the scope of the user's request. +- Do not refactor, optimize, or rewrite unless explicitly instructed. +- Ensure the edited code remains complete and executable. +- Output only the modified code, without any additional explanation, comments, or extra formatting. + +Make your edits precise, minimal, and faithful to the user's instructions, using the dataset context to guide your modifications. + """ + dataset_context = dspy.InputField(desc="The dataset context to be used for the code edit, including information about the dataset's shape, columns, types, and null values") + original_code = dspy.InputField(desc="The original code the user wants modified") + user_prompt = dspy.InputField(desc="The user instruction describing how the code should be changed") + edited_code = dspy.OutputField(desc="The updated version of the code reflecting the user's request, incorporating changes informed by the dataset context") + +# The ind module is called when agent_name is +# explicitly mentioned in the query +class auto_analyst_ind(dspy.Module): + """Handles individual agent execution when explicitly specified in query""" + + def __init__(self, agents, retrievers): + # Initialize agent modules and retrievers + self.agents = {} + self.agent_inputs = {} + self.agent_desc = [] + + # Create modules from agent signatures + for i, a in enumerate(agents): + name = a.__pydantic_core_schema__['schema']['model_name'] + self.agents[name] = dspy.ChainOfThoughtWithHint(a) + self.agent_inputs[name] = {x.strip() for x in str(agents[i].__pydantic_core_schema__['cls']).split('->')[0].split('(')[1].split(',')} + self.agent_desc.append(get_agent_description(name)) + + # Initialize components + self.memory_summarize_agent = dspy.ChainOfThought(m.memory_summarize_agent) + self.dataset = retrievers['dataframe_index'].as_retriever(k=1) + self.styling_index = retrievers['style_index'].as_retriever(similarity_top_k=1) + self.code_combiner_agent = dspy.ChainOfThought(code_combiner_agent) + + # Initialize thread pool + self.executor = ThreadPoolExecutor(max_workers=min(4, os.cpu_count() * 2)) + + def execute_agent(self, specified_agent, inputs): + """Execute agent and generate memory summary in parallel""" + try: + # Execute main agent + agent_result = self.agents[specified_agent.strip()](**inputs) + return specified_agent.strip(), dict(agent_result) + + except Exception as e: + return specified_agent.strip(), {"error": str(e)} + + def execute_agent_with_memory(self, specified_agent, inputs, query): + """Execute agent and generate memory summary in parallel""" + try: + # Execute main agent + agent_result = self.agents[specified_agent.strip()](**inputs) + agent_dict = dict(agent_result) + + # Generate memory summary + memory_result = self.memory_summarize_agent( + agent_response=specified_agent+' '+agent_dict['code']+'\n'+agent_dict['summary'], + user_goal=query + ) + + return { + specified_agent.strip(): agent_dict, + 'memory_'+specified_agent.strip(): str(memory_result.summary) + } + except Exception as e: + return {"error": str(e)} + + def forward(self, query, specified_agent): + try: + # If specified_agent contains multiple agents separated by commas + # This is for handling multiple @agent mentions in one query + if "," in specified_agent: + agent_list = [agent.strip() for agent in specified_agent.split(",")] + return self.execute_multiple_agents(query, agent_list) + + # Process query with specified agent (single agent case) + dict_ = {} + dict_['dataset'] = self.dataset.retrieve(query)[0].text + dict_['styling_index'] = self.styling_index.retrieve(query)[0].text + dict_['hint'] = [] + dict_['goal'] = query + dict_['Agent_desc'] = str(self.agent_desc) + + # Prepare inputs + inputs = {x:dict_[x] for x in self.agent_inputs[specified_agent.strip()]} + inputs['hint'] = str(dict_['hint']).replace('[','').replace(']','') + + # Execute agent + result = self.agents[specified_agent.strip()](**inputs) + output_dict = {specified_agent.strip(): dict(result)} + + if "error" in output_dict: + return {"response": f"Error executing agent: {output_dict['error']}"} + + return output_dict + + except Exception as e: + return {"response": f"This is the error from the system: {str(e)}"} + + def execute_multiple_agents(self, query, agent_list): + """Execute multiple agents sequentially on the same query""" + try: + # Initialize resources + dict_ = {} + dict_['dataset'] = self.dataset.retrieve(query)[0].text + dict_['styling_index'] = self.styling_index.retrieve(query)[0].text + dict_['hint'] = [] + dict_['goal'] = query + dict_['Agent_desc'] = str(self.agent_desc) + + results = {} + code_list = [] + + # Execute each agent sequentially + for agent_name in agent_list: + if agent_name not in self.agents: + results[agent_name] = {"error": f"Agent '{agent_name}' not found"} + continue + + # Prepare inputs for this agent + inputs = {x:dict_[x] for x in self.agent_inputs[agent_name] if x in dict_} + inputs['hint'] = str(dict_['hint']).replace('[','').replace(']','') + + # Execute agent + agent_result = self.agents[agent_name](**inputs) + agent_dict = dict(agent_result) + results[agent_name] = agent_dict + + # Collect code for later combination + if 'code' in agent_dict: + code_list.append(agent_dict['code']) + + return results + + except Exception as e: + return {"response": f"Error executing multiple agents: {str(e)}"} + + +# This is the auto_analyst with planner +class auto_analyst(dspy.Module): + """Main analyst module that coordinates multiple agents using a planner""" + + def __init__(self, agents, retrievers): + # Initialize agent modules and retrievers + self.agents = {} + self.agent_inputs = {} + self.agent_desc = [] + + for i, a in enumerate(agents): + name = a.__pydantic_core_schema__['schema']['model_name'] + self.agents[name] = dspy.ChainOfThought(a) + self.agent_inputs[name] = {x.strip() for x in str(agents[i].__pydantic_core_schema__['cls']).split('->')[0].split('(')[1].split(',')} + self.agent_desc.append({name: get_agent_description(name)}) + + # Initialize coordination agents + self.planner = dspy.ChainOfThought(analytical_planner) + self.refine_goal = dspy.ChainOfThought(goal_refiner_agent) + self.code_combiner_agent = dspy.ChainOfThought(code_combiner_agent) + self.story_teller = dspy.ChainOfThought(story_teller_agent) + self.memory_summarize_agent = dspy.ChainOfThought(m.memory_summarize_agent) + + # Initialize retrievers + self.dataset = retrievers['dataframe_index'].as_retriever(k=1) + self.styling_index = retrievers['style_index'].as_retriever(similarity_top_k=1) + + # Initialize thread pool for parallel execution + self.executor = ThreadPoolExecutor(max_workers=min(len(agents) + 2, os.cpu_count() * 2)) + + def execute_agent(self, agent_name, inputs): + """Execute a single agent with given inputs""" + try: + result = self.agents[agent_name.strip()](**inputs) + return agent_name.strip(), dict(result) + except Exception as e: + return agent_name.strip(), {"error": str(e)} + + def get_plan(self, query): + """Get the analysis plan""" + dict_ = {} + dict_['dataset'] = self.dataset.retrieve(query)[0].text + dict_['styling_index'] = self.styling_index.retrieve(query)[0].text + dict_['goal'] = query + dict_['Agent_desc'] = str(self.agent_desc) + + plan = self.planner(goal=dict_['goal'], dataset=dict_['dataset'], Agent_desc=dict_['Agent_desc']) + return dict(plan) + + async def execute_plan(self, query, plan): + """Execute the plan and yield results as they complete""" + dict_ = {} + dict_['dataset'] = self.dataset.retrieve(query)[0].text + dict_['styling_index'] = self.styling_index.retrieve(query)[0].text + dict_['hint'] = [] + dict_['goal'] = query + + import json + + # Clean and split the plan string into agent names + plan_text = plan.plan.replace("Plan", "").replace(":", "").strip() + plan_list = [agent.strip() for agent in plan_text.split("->") if agent.strip()] + # logger.log(f"Plan list: {plan_list}") + # Parse the attached plan_instructions into a dict + raw_instr = plan.plan_instructions + # logger.log(f"Raw instructions: {raw_instr}") + if isinstance(raw_instr, str): + try: + plan_instructions = json.loads(raw_instr) + except Exception: + plan_instructions = {} + elif isinstance(raw_instr, dict): + plan_instructions = raw_instr + else: + plan_instructions = {} + logger + # If no plan was produced, short-circuit + if not plan_list: + yield "plan_not_found", dict(plan), {"error": "No plan found"} + return + + # Launch each agent in parallel, attaching its own instructions + futures = [] + for agent_name in plan_list: + key = agent_name.strip() + # gather input fields except plan_instructions + inputs = { + param: dict_[param] + for param in self.agent_inputs[key] + if param != "plan_instructions" + } + # attach the specific instructions for this agent (or defaults) + if "plan_instructions" in self.agent_inputs[key]: + inputs['plan_instructions'] = plan_instructions + inputs["your_task"] = str(plan_instructions.get( + key, "" + )) + # logger.log(f"Inputs: {inputs}") + future = self.executor.submit(self.execute_agent, agent_name, inputs) + futures.append((agent_name, inputs, future)) + # Yield results as they complete + completed_results = [] + for agent_name, inputs, future in futures: + try: + name, result = await asyncio.get_event_loop().run_in_executor(None, future.result) + completed_results.append((name, result)) + yield name, inputs, result + except Exception as e: + yield agent_name, inputs, {"error": str(e)} diff --git a/src/db/schemas/models.py b/src/db/schemas/models.py index c26d6084e30323c2538fb1308e33fd9700c7a5d8..1087c85f04438a98c19af7d21c2796e3ce8887df 100644 --- a/src/db/schemas/models.py +++ b/src/db/schemas/models.py @@ -1,7 +1,7 @@ -from sqlalchemy import create_engine, Column, Integer, String, ForeignKey, DateTime, Text, Float, Boolean, JSON, UniqueConstraint +from sqlalchemy import create_engine, Column, Integer, String, ForeignKey, DateTime, Text, Float, Boolean from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker, relationship -from datetime import datetime, UTC +from datetime import datetime # Define the base class for declarative models Base = declarative_base() @@ -11,14 +11,12 @@ class User(Base): __tablename__ = 'users' user_id = Column(Integer, primary_key=True, autoincrement=True) - username = Column(String, nullable=False) # Removed unique=True + username = Column(String, unique=True, nullable=False) email = Column(String, unique=True, nullable=False) - created_at = Column(DateTime, default=lambda: datetime.now(UTC)) + created_at = Column(DateTime, default=datetime.utcnow) # Add relationship for cascade options chats = relationship("Chat", back_populates="user", cascade="all, delete-orphan") usage_records = relationship("ModelUsage", back_populates="user") - deep_analysis_reports = relationship("DeepAnalysisReport", back_populates="user", cascade="all, delete-orphan") - template_preferences = relationship("UserTemplatePreference", back_populates="user", cascade="all, delete-orphan") # Define the Chats table class Chat(Base): @@ -27,7 +25,7 @@ class Chat(Base): chat_id = Column(Integer, primary_key=True, autoincrement=True) user_id = Column(Integer, ForeignKey('users.user_id', ondelete="CASCADE"), nullable=True) title = Column(String, default='New Chat') - created_at = Column(DateTime, default=lambda: datetime.now(UTC)) + created_at = Column(DateTime, default=datetime.utcnow) # Add relationships for cascade options user = relationship("User", back_populates="chats") messages = relationship("Message", back_populates="chat", cascade="all, delete-orphan") @@ -41,10 +39,9 @@ class Message(Base): chat_id = Column(Integer, ForeignKey('chats.chat_id', ondelete="CASCADE"), nullable=False) sender = Column(String, nullable=False) # 'user' or 'ai' content = Column(Text, nullable=False) - timestamp = Column(DateTime, default=lambda: datetime.now(UTC)) + timestamp = Column(DateTime, default=datetime.utcnow) # Add relationship for cascade options chat = relationship("Chat", back_populates="messages") - feedback = relationship("MessageFeedback", back_populates="message", uselist=False, cascade="all, delete-orphan") # Define the Model Usage table class ModelUsage(Base): @@ -62,176 +59,9 @@ class ModelUsage(Base): query_size = Column(Integer, default=0) # Size in characters response_size = Column(Integer, default=0) # Size in characters cost = Column(Float, default=0.0) # Cost in USD - timestamp = Column(DateTime, default=lambda: datetime.now(UTC)) + timestamp = Column(DateTime, default=datetime.utcnow) is_streaming = Column(Boolean, default=False) request_time_ms = Column(Integer, default=0) # Request processing time in milliseconds # Add relationships user = relationship("User", back_populates="usage_records") - chat = relationship("Chat", back_populates="usage_records") - -# Define the Code Execution table -class CodeExecution(Base): - """Tracks code execution attempts and results for analysis and debugging.""" - __tablename__ = 'code_executions' - - execution_id = Column(Integer, primary_key=True, autoincrement=True) - message_id = Column(Integer, ForeignKey('messages.message_id', ondelete="CASCADE"), nullable=True) - chat_id = Column(Integer, ForeignKey('chats.chat_id', ondelete="CASCADE"), nullable=True) - user_id = Column(Integer, ForeignKey('users.user_id', ondelete="SET NULL"), nullable=True) - - # Code tracking - initial_code = Column(Text, nullable=True) # First version of code submitted - latest_code = Column(Text, nullable=True) # Most recent version of code - - # Execution results - is_successful = Column(Boolean, default=False) - output = Column(Text, nullable=True) # Full output including errors - - # Model and agent information - model_provider = Column(String(50), nullable=True) - model_name = Column(String(100), nullable=True) - model_temperature = Column(Float, nullable=True) - model_max_tokens = Column(Integer, nullable=True) - - # Failure information - failed_agents = Column(Text, nullable=True) # JSON list of agent names that failed - error_messages = Column(Text, nullable=True) # JSON map of error messages by agent - - # Metadata - created_at = Column(DateTime, default=lambda: datetime.now(UTC)) - updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) - -class MessageFeedback(Base): - """Tracks user feedback and model settings for each message.""" - __tablename__ = 'message_feedback' - - feedback_id = Column(Integer, primary_key=True, autoincrement=True) - message_id = Column(Integer, ForeignKey('messages.message_id', ondelete="CASCADE"), nullable=False) - - # User feedback - rating = Column(Integer, nullable=True) # Star rating (1-5) - - # Model settings used for this message - model_name = Column(String(100), nullable=True) - model_provider = Column(String(50), nullable=True) - temperature = Column(Float, nullable=True) - max_tokens = Column(Integer, nullable=True) - - # Metadata - created_at = Column(DateTime, default=lambda: datetime.now(UTC)) - updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) - - # Relationship - message = relationship("Message", back_populates="feedback") - -class DeepAnalysisReport(Base): - """Stores deep analysis reports with comprehensive analysis data and metadata.""" - __tablename__ = 'deep_analysis_reports' - - report_id = Column(Integer, primary_key=True, autoincrement=True) - report_uuid = Column(String(100), unique=True, nullable=False) # Frontend generated ID - user_id = Column(Integer, ForeignKey('users.user_id', ondelete="CASCADE"), nullable=True) - - # Analysis objective and status - goal = Column(Text, nullable=False) # The analysis objective/question - status = Column(String(20), nullable=False, default='pending') # 'pending', 'running', 'completed', 'failed' - - # Timing information - start_time = Column(DateTime, default=lambda: datetime.now(UTC)) - end_time = Column(DateTime, nullable=True) - duration_seconds = Column(Integer, nullable=True) # Calculated duration - - # Analysis components (stored as text/JSON) - deep_questions = Column(Text, nullable=True) # Generated analytical questions - deep_plan = Column(Text, nullable=True) # Analysis plan - summaries = Column(JSON, nullable=True) # Array of analysis summaries - analysis_code = Column(Text, nullable=True) # Generated Python code - plotly_figures = Column(JSON, nullable=True) # Array of Plotly figure data - synthesis = Column(JSON, nullable=True) # Array of synthesis insights - final_conclusion = Column(Text, nullable=True) # Final analysis conclusion - - # Report output - html_report = Column(Text, nullable=True) # Complete HTML report - report_summary = Column(Text, nullable=True) # Brief summary for listing - - # Execution tracking - progress_percentage = Column(Integer, default=0) # Progress 0-100 - steps_completed = Column(JSON, nullable=True) # Array of completed step names - error_message = Column(Text, nullable=True) # Error details if failed - - # Model and cost tracking - model_provider = Column(String(50), nullable=True) - model_name = Column(String(100), nullable=True) - total_tokens_used = Column(Integer, default=0) - estimated_cost = Column(Float, default=0.0) # Cost in USD - credits_consumed = Column(Integer, default=0) # Credits deducted for this analysis - - # Metadata - created_at = Column(DateTime, default=lambda: datetime.now(UTC)) - updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) - - # Relationships - user = relationship("User", back_populates="deep_analysis_reports") - -class AgentTemplate(Base): - """Stores predefined agent templates that users can enable/disable.""" - __tablename__ = 'agent_templates' - - template_id = Column(Integer, primary_key=True, autoincrement=True) - - # Template definition - template_name = Column(String(100), nullable=False, unique=True) # e.g., 'pytorch_specialist', 'data_cleaning_expert' - display_name = Column(String(200), nullable=True) # User-friendly display name - description = Column(Text, nullable=False) # Short description for template selection - prompt_template = Column(Text, nullable=False) # Main prompt/instructions for agent behavior - - # Template appearance - icon_url = Column(String(500), nullable=True) # URL to template icon (CDN, data URL, or relative path) - - # Template categorization - category = Column(String(50), nullable=True) # 'Visualization', 'Modelling', 'Data Manipulation' - is_premium_only = Column(Boolean, default=False) # True if template requires premium subscription - - # Agent variant support - variant_type = Column(String(20), default='individual') # 'planner', 'individual', or 'both' - base_agent = Column(String(100), nullable=True) # Base agent name for variants (e.g., 'preprocessing_agent') - - # Status and metadata - is_active = Column(Boolean, default=True) - - # Timestamps - created_at = Column(DateTime, default=lambda: datetime.now(UTC)) - updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) - - # Relationships - user_preferences = relationship("UserTemplatePreference", back_populates="template", cascade="all, delete-orphan") - -class UserTemplatePreference(Base): - """Tracks user preferences and usage for agent templates.""" - __tablename__ = 'user_template_preferences' - - preference_id = Column(Integer, primary_key=True, autoincrement=True) - user_id = Column(Integer, ForeignKey('users.user_id', ondelete="CASCADE"), nullable=False) - template_id = Column(Integer, ForeignKey('agent_templates.template_id', ondelete="CASCADE"), nullable=False) - - # User preferences - is_enabled = Column(Boolean, default=True) # Whether user has this template enabled - - # Usage tracking - usage_count = Column(Integer, default=0) # Track how many times user has used this template - last_used_at = Column(DateTime, nullable=True) # Last time user used this template - - # Timestamps - created_at = Column(DateTime, default=lambda: datetime.now(UTC)) - updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) - - # Relationships - user = relationship("User", back_populates="template_preferences") - template = relationship("AgentTemplate", back_populates="user_preferences") - - # Constraints - user can only have one preference record per template - __table_args__ = ( - UniqueConstraint('user_id', 'template_id', name='unique_user_template_preference'), - ) - - \ No newline at end of file + chat = relationship("Chat", back_populates="usage_records") \ No newline at end of file diff --git a/src/managers/ai_manager.py b/src/managers/ai_manager.py index 6446ba296adba94d0deb8f96cd499ab02f635edb..b5e10f82854f8fd2c97c743b87602184d0b8ab42 100644 --- a/src/managers/ai_manager.py +++ b/src/managers/ai_manager.py @@ -1,8 +1,10 @@ import logging from typing import Optional, Dict, Any +import time from src.db.schemas.models import ModelUsage from src.db.init_db import session_factory -from datetime import datetime, UTC +from datetime import datetime +import tiktoken from src.routes.analytics_routes import handle_new_model_usage import asyncio @@ -43,7 +45,7 @@ class AI_Manager: query_size=query_size, response_size=response_size, cost=cost, - timestamp=datetime.now(UTC), + timestamp=datetime.utcnow(), is_streaming=is_streaming, request_time_ms=request_time_ms ) diff --git a/src/managers/app_manager.py b/src/managers/app_manager.py deleted file mode 100644 index 63fe2e230d323c98f1a99d07569a010871debe05..0000000000000000000000000000000000000000 --- a/src/managers/app_manager.py +++ /dev/null @@ -1,127 +0,0 @@ -import logging -import dspy -from src.managers.session_manager import SessionManager -from src.managers.ai_manager import AI_Manager -from src.utils.logger import Logger - -logger = Logger("app_manager", see_time=True, console_log=False) - -class AppState: - def __init__(self, styling_instructions, chat_history_name_agent, default_model_config): - self._session_manager = SessionManager(styling_instructions, {}) # Empty dict, agents loaded from DB - self.model_config = default_model_config.copy() - - # Update the SessionManager with the current model_config - self._session_manager._app_model_config = self.model_config - - self.ai_manager = AI_Manager() - self.chat_name_agent = chat_history_name_agent - - # Initialize deep analysis module - self.deep_analyzer = None - - def get_session_state(self, session_id: str): - """Get or create session-specific state using the SessionManager""" - return self._session_manager.get_session_state(session_id) - - def clear_session_state(self, session_id: str): - """Clear session-specific state using the SessionManager""" - self._session_manager.clear_session_state(session_id) - - def update_session_dataset(self, session_id: str, datasets, names, desc, pre_generated=False): - """Update dataset for a specific session using the SessionManager""" - self._session_manager.update_session_dataset(session_id, datasets, names, desc, pre_generated=pre_generated) - - def reset_session_to_default(self, session_id: str): - """Reset a session to use the default dataset using the SessionManager""" - self._session_manager.reset_session_to_default(session_id) - - def set_session_user(self, session_id: str, user_id: int, chat_id: int = None): - """Associate a user with a session using the SessionManager""" - return self._session_manager.set_session_user(session_id, user_id, chat_id) - - def get_ai_manager(self): - """Get the AI Manager instance""" - return self.ai_manager - - def get_provider_for_model(self, model_name): - return self.ai_manager.get_provider_for_model(model_name) - - def calculate_cost(self, model_name, input_tokens, output_tokens): - return self.ai_manager.calculate_cost(model_name, input_tokens, output_tokens) - - def save_usage_to_db(self, user_id, chat_id, model_name, provider, prompt_tokens, completion_tokens, total_tokens, query_size, response_size, cost, request_time_ms, is_streaming=False): - return self.ai_manager.save_usage_to_db(user_id, chat_id, model_name, provider, prompt_tokens, completion_tokens, total_tokens, query_size, response_size, round(cost, 7), request_time_ms, is_streaming) - - def get_tokenizer(self): - return self.ai_manager.tokenizer - - def get_chat_history_name_agent(self): - return dspy.Predict(self.chat_name_agent) - - def get_deep_analyzer(self, session_id: str): - """Get or create deep analysis module for a session""" - session_state = self.get_session_state(session_id) - user_id = session_state.get("user_id") - - # Check if we need to recreate the deep analyzer (user changed or doesn't exist) - current_analyzer = session_state.get('deep_analyzer') - analyzer_user_id = session_state.get('deep_analyzer_user_id') - - logger.log_message(f"Deep analyzer check - session: {session_id}, current_user: {user_id}, analyzer_user: {analyzer_user_id}, has_analyzer: {current_analyzer is not None}", level=logging.INFO) - - if (not current_analyzer or - analyzer_user_id != user_id or - not hasattr(session_state, 'deep_analyzer')): - - logger.log_message(f"Creating/recreating deep analyzer for session {session_id}, user_id: {user_id} (reason: analyzer_exists={current_analyzer is not None}, user_match={analyzer_user_id == user_id})", level=logging.INFO) - - # Load user-enabled agents from database using preference system - from src.db.init_db import session_factory - from src.agents.agents import load_user_enabled_templates_for_planner_from_db - - db_session = session_factory() - try: - # Load user-enabled agents for planner (respects preferences) - if user_id: - enabled_agents_dict = load_user_enabled_templates_for_planner_from_db(user_id, db_session) - else: - enabled_agents_dict = {} - - if not enabled_agents_dict: - # Fallback to default agents if no user preferences - enabled_agents_dict = { - "preprocessing_agent": "preprocessing_agent", - "statistical_analytics_agent": "statistical_analytics_agent", - "sk_learn_agent": "sk_learn_agent", - "data_viz_agent": "data_viz_agent" - } - - # Import deep analysis module - from src.agents.deep_agents import deep_analysis_module - from src.agents.agents import get_agent_description - - deep_agents = {} - deep_agents_desc = {} - - for agent_name, signature in enabled_agents_dict.items(): - deep_agents[agent_name] = signature - deep_agents_desc[agent_name] = get_agent_description(agent_name) - - logger.log_message(f"Deep analyzer initialized with {len(deep_agents)} agents: {list(deep_agents.keys())}", level=logging.INFO) - - finally: - db_session.close() - - session_state['deep_analyzer'] = deep_analysis_module( - agents=deep_agents, - agents_desc=deep_agents_desc - ) - # Set datasets separately or pass them when needed - session_state['deep_analyzer'].datasets = session_state.get("datasets") - session_state['deep_analyzer_user_id'] = user_id # Track which user this analyzer was created for - - else: - logger.log_message(f"Using existing deep analyzer for session {session_id}, user_id: {user_id}", level=logging.INFO) - - return session_state['deep_analyzer'] diff --git a/src/managers/chat_manager.py b/src/managers/chat_manager.py index 7038f69e76195a2f2c1e0da429d84411059a1327..04de71b4f98d1e85193c793b72de76936f24dafc 100644 --- a/src/managers/chat_manager.py +++ b/src/managers/chat_manager.py @@ -1,12 +1,16 @@ -from sqlalchemy import create_engine, func, exists + +from sqlalchemy import create_engine, desc, func, exists from sqlalchemy.orm import sessionmaker, scoped_session from sqlalchemy.exc import SQLAlchemyError -from src.db.schemas.models import Base, User, Chat, Message, ModelUsage, MessageFeedback +from src.db.schemas.models import Base, User, Chat, Message, ModelUsage import logging -from typing import List, Dict, Optional, Any -from datetime import datetime, UTC +import requests +import json +from typing import List, Dict, Optional, Tuple, Any +from datetime import datetime +import time +import tiktoken from src.utils.logger import Logger -from src.utils.model_registry import MODEL_COSTS import re logger = Logger("chat_manager", see_time=True, console_log=False) @@ -30,7 +34,39 @@ class ChatManager: self.Session = scoped_session(sessionmaker(bind=self.engine)) # Add price mappings for different models - self.model_costs = MODEL_COSTS + self.model_costs = { + "openai": { + "gpt-4": {"input": 0.03, "output": 0.06}, + "gpt-4o": {"input": 0.0025, "output": 0.01}, + "gpt-4.5-preview": {"input": 0.075, "output": 0.15}, + "gpt-4o-mini": {"input": 0.00015, "output": 0.0006}, + "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015}, + "o1": {"input": 0.015, "output": 0.06}, + "o1-mini": {"input": 0.00011, "output": 0.00044}, + "o3-mini": {"input": 0.00011, "output": 0.00044} + }, + "anthropic": { + "claude-3-opus-latest": {"input": 0.015, "output": 0.075}, + "claude-3-7-sonnet-latest": {"input": 0.003, "output": 0.015}, + "claude-3-5-sonnet-latest": {"input": 0.003, "output": 0.015}, + "claude-3-5-haiku-latest": {"input": 0.0008, "output": 0.0004}, + }, + "groq": { + "deepseek-r1-distill-llama-70b": {"input": 0.00075, "output": 0.00099}, + "llama-3.3-70b-versatile": {"input": 0.00059, "output": 0.00079}, + "llama3-8b-8192": {"input": 0.00005, "output": 0.00008}, + "llama3-70b-8192": {"input": 0.00059, "output": 0.00079}, + "llama-3.1-8b-instant": {"input": 0.00005, "output": 0.00008}, + "mistral-saba-24b": {"input": 0.00079, "output": 0.00079}, + "gemma2-9b-it": {"input": 0.0002, "output": 0.0002}, + "qwen-qwq-32b": {"input": 0.00029, "output": 0.00039}, + "meta-llama/llama-4-maverick-17b-128e-instruct": {"input": 0.0002, "output": 0.0006}, + "meta-llama/llama-4-scout-17b-16e-instruct": {"input": 0.00011, "output": 0.00034}, + }, + "gemini": { + "gemini-2.5-pro-preview-03-25": {"input": 0.00015, "output": 0.001} + } + } # Add model providers mapping @@ -57,7 +93,7 @@ class ChatManager: chat = Chat( user_id=user_id, title='New Chat', - created_at=datetime.now(UTC) + created_at=datetime.utcnow() ) session.add(chat) session.flush() # Flush to get the ID before commit @@ -116,7 +152,7 @@ class ChatManager: chat_id=chat_id, content=content, sender=sender, - timestamp=datetime.now(UTC) + timestamp=datetime.utcnow() ) session.add(message) session.flush() # Flush to get the ID before commit @@ -699,246 +735,3 @@ class ChatManager: except Exception as e: logger.log_message(f"Error in extract_response_history: {str(e)}", level=logging.ERROR) return "" - - def add_message_feedback(self, message_id: int, rating: int, - model_settings: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: - """ - Add or update feedback for a message. - - Args: - message_id: ID of the message to add feedback for - rating: Star rating (1-5) - model_settings: Optional dictionary containing model settings (name, provider, temperature, etc.) - - Returns: - Dictionary containing feedback information - """ - session = self.Session() - try: - # Check if message exists - message = session.query(Message).filter(Message.message_id == message_id).first() - if not message: - raise ValueError(f"Message with ID {message_id} not found") - - # Check if feedback already exists - existing_feedback = session.query(MessageFeedback).filter( - MessageFeedback.message_id == message_id - ).first() - - now = datetime.now(UTC) - - # Extract model settings - model_name = None - model_provider = None - temperature = None - max_tokens = None - - if model_settings: - model_name = model_settings.get('model_name') - model_provider = model_settings.get('model_provider') - temperature = model_settings.get('temperature') - max_tokens = model_settings.get('max_tokens') - - if existing_feedback: - # Update existing feedback - existing_feedback.rating = rating - existing_feedback.model_name = model_name - existing_feedback.model_provider = model_provider - existing_feedback.temperature = temperature - existing_feedback.max_tokens = max_tokens - existing_feedback.updated_at = now - feedback_record = existing_feedback - else: - # Create new feedback - feedback_record = MessageFeedback( - message_id=message_id, - rating=rating, - model_name=model_name, - model_provider=model_provider, - temperature=temperature, - max_tokens=max_tokens, - created_at=now, - updated_at=now - ) - session.add(feedback_record) - - session.commit() - - return { - "feedback_id": feedback_record.feedback_id, - "message_id": feedback_record.message_id, - "rating": feedback_record.rating, - "model_name": feedback_record.model_name, - "model_provider": feedback_record.model_provider, - "temperature": feedback_record.temperature, - "max_tokens": feedback_record.max_tokens, - "created_at": feedback_record.created_at.isoformat(), - "updated_at": feedback_record.updated_at.isoformat() - } - except SQLAlchemyError as e: - session.rollback() - logger.log_message(f"Error adding feedback: {str(e)}", level=logging.ERROR) - raise - finally: - session.close() - - def get_message_feedback(self, message_id: int) -> Optional[Dict[str, Any]]: - """ - Get feedback for a specific message. - - Args: - message_id: ID of the message to get feedback for - - Returns: - Dictionary containing feedback information or None if no feedback exists - """ - session = self.Session() - try: - feedback = session.query(MessageFeedback).filter( - MessageFeedback.message_id == message_id - ).first() - - if not feedback: - return None - - return { - "feedback_id": feedback.feedback_id, - "message_id": feedback.message_id, - "rating": feedback.rating, - "model_name": feedback.model_name, - "model_provider": feedback.model_provider, - "temperature": feedback.temperature, - "max_tokens": feedback.max_tokens, - "created_at": feedback.created_at.isoformat(), - "updated_at": feedback.updated_at.isoformat() - } - except SQLAlchemyError as e: - logger.log_message(f"Error getting feedback: {str(e)}", level=logging.ERROR) - raise - finally: - session.close() - - def get_chat_feedback(self, chat_id: int) -> List[Dict[str, Any]]: - """ - Get all feedback for messages in a specific chat. - - Args: - chat_id: ID of the chat to get feedback for - - Returns: - List of dictionaries containing feedback information - """ - session = self.Session() - try: - feedback_records = session.query(MessageFeedback).join( - Message, Message.message_id == MessageFeedback.message_id - ).filter( - Message.chat_id == chat_id - ).all() - - return [{ - "feedback_id": feedback.feedback_id, - "message_id": feedback.message_id, - "rating": feedback.rating, - "model_name": feedback.model_name, - "model_provider": feedback.model_provider, - "temperature": feedback.temperature, - "max_tokens": feedback.max_tokens, - "created_at": feedback.created_at.isoformat(), - "updated_at": feedback.updated_at.isoformat() - } for feedback in feedback_records] - except SQLAlchemyError as e: - logger.log_message(f"Error getting chat feedback: {str(e)}", level=logging.ERROR) - raise - finally: - session.close() - - def get_feedback_statistics(self, user_id: Optional[int] = None, - start_date: Optional[datetime] = None, - end_date: Optional[datetime] = None) -> Dict[str, Any]: - """ - Get feedback statistics for analysis. - - Args: - user_id: Optional user ID to filter by - start_date: Optional start date to filter by - end_date: Optional end date to filter by - - Returns: - Dictionary containing feedback statistics - """ - session = self.Session() - try: - # Base query for all feedback - query = session.query(MessageFeedback).join( - Message, Message.message_id == MessageFeedback.message_id - ) - - # Apply filters if provided - if user_id is not None: - query = query.join(Chat, Chat.chat_id == Message.chat_id).filter( - Chat.user_id == user_id - ) - - if start_date is not None: - query = query.filter(MessageFeedback.created_at >= start_date) - - if end_date is not None: - query = query.filter(MessageFeedback.created_at <= end_date) - - # Get all feedback records - feedback_records = query.all() - - # Calculate statistics - if not feedback_records: - return { - "total_feedback_count": 0, - "average_rating": 0, - "rating_distribution": { - "1": 0, "2": 0, "3": 0, "4": 0, "5": 0 - }, - "model_ratings": {} - } - - # Calculate average rating - ratings = [record.rating for record in feedback_records if record.rating is not None] - average_rating = sum(ratings) / len(ratings) if ratings else 0 - - # Calculate rating distribution - rating_distribution = { - "1": 0, "2": 0, "3": 0, "4": 0, "5": 0 - } - - for record in feedback_records: - if record.rating is not None: - rating_distribution[str(record.rating)] += 1 - - # Calculate ratings by model - model_ratings = {} - for record in feedback_records: - if record.model_name and record.rating is not None: - if record.model_name not in model_ratings: - model_ratings[record.model_name] = { - "count": 0, - "total": 0, - "average": 0 - } - - model_ratings[record.model_name]["count"] += 1 - model_ratings[record.model_name]["total"] += record.rating - - # Calculate average for each model - for model_name, data in model_ratings.items(): - data["average"] = data["total"] / data["count"] if data["count"] > 0 else 0 - - return { - "total_feedback_count": len(feedback_records), - "average_rating": average_rating, - "rating_distribution": rating_distribution, - "model_ratings": model_ratings - } - except SQLAlchemyError as e: - logger.log_message(f"Error getting feedback statistics: {str(e)}", level=logging.ERROR) - raise - finally: - session.close() diff --git a/src/managers/session_manager.py b/src/managers/session_manager.py index d73807143cb73147ccc7a7b7358f5d956ae37242..b5724ecfc9667c429b70e9a7e4e8d1715e4674c0 100644 --- a/src/managers/session_manager.py +++ b/src/managers/session_manager.py @@ -4,31 +4,21 @@ import time import uuid import logging import pandas as pd -from typing import Dict, Any, List +from typing import Dict, Any, List, Optional -from fastapi import HTTPException -from src.utils.simple_retriever import Document, SimpleRetriever +from llama_index.core import Document, VectorStoreIndex from src.utils.logger import Logger -from src.managers.user_manager import get_current_user -from src.agents.agents import auto_analyst, dataset_description_agent, data_context_gen +from src.managers.user_manager import create_user, get_current_user +from src.agents.agents import auto_analyst, auto_analyst_ind from src.agents.retrievers.retrievers import make_data from src.managers.chat_manager import ChatManager -from src.utils.model_registry import mid_lm from dotenv import load_dotenv -import duckdb -import dspy -from src.utils.dataset_description_generator import generate_dataset_description -from fastapi import Request load_dotenv() # Initialize logger logger = Logger("session_manager", see_time=False, console_log=False) -# Helper to clamp temperature to valid range -def _get_clamped_temperature(): - return min(1.0, max(0.0, float(os.getenv("TEMPERATURE", "1.0")))) - class SessionManager: """ Manages session-specific state, including datasets, retrievers, and AI systems. @@ -37,23 +27,19 @@ class SessionManager: def __init__(self, styling_instructions: List[str], available_agents: Dict): """ - Initialize SessionManager with styling instructions and available agents + Initialize session manager with styling instructions and agents Args: - styling_instructions: List of styling instructions for visualization - available_agents: Dictionary of available agents (deprecated - agents now loaded from DB) + styling_instructions: List of styling instructions + available_agents: Dictionary of available agents """ - self.styling_instructions = styling_instructions self._sessions = {} self._default_df = None self._default_retrievers = None self._default_ai_system = None + self._dataset_description = None self._make_data = None - - # Initialize chat manager - - self._default_name = "Housing.csv" - + self._default_name = "Housing Dataset" # Default dataset name self._dataset_description = """This dataset contains residential property information with details about pricing, physical characteristics, and amenities. The data can be used for real estate market analysis, property valuation, and understanding the relationship between house features and prices. @@ -94,6 +80,7 @@ Data Handling Recommendations: This dataset appears clean with consistent formatting and no missing values, making it suitable for immediate analysis with appropriate categorical encoding. """ + self.styling_instructions = styling_instructions self.available_agents = available_agents self.chat_manager = ChatManager(db_url=os.getenv("DATABASE_URL")) @@ -103,19 +90,29 @@ This dataset appears clean with consistent formatting and no missing values, mak """Initialize the default dataset and store it""" try: self._default_df = pd.read_csv("Housing.csv") - self._make_data = {'dataset_python_name':"this dataset is loaded as `df`","description":self._dataset_description} + self._make_data = make_data(self._default_df, self._dataset_description) self._default_retrievers = self.initialize_retrievers(self.styling_instructions, [str(self._make_data)]) - # Create default AI system - agents will be loaded from database - self._default_ai_system = auto_analyst(agents=[], retrievers=self._default_retrievers) + self._default_ai_system = auto_analyst(agents=list(self.available_agents.values()), + retrievers=self._default_retrievers) except Exception as e: logger.log_message(f"Error initializing default dataset: {str(e)}", level=logging.ERROR) raise e - def initialize_retrievers(self,styling_instructions: List[str], doc: List[str]): - try: - style_index = SimpleRetriever.from_documents([Document(text=x) for x in styling_instructions]) + def initialize_retrievers(self, styling_instructions: List[str], doc: List[str]): + """ + Initialize retrievers for styling and data + + Args: + styling_instructions: List of styling instructions + doc: List of document strings - return {"style_index": style_index, "dataframe_index": doc} + Returns: + Dictionary containing style_index and dataframe_index + """ + try: + style_index = VectorStoreIndex.from_documents([Document(text=x) for x in styling_instructions]) + data_index = VectorStoreIndex.from_documents([Document(text=x) for x in doc]) + return {"style_index": style_index, "dataframe_index": data_index} except Exception as e: logger.log_message(f"Error initializing retrievers: {str(e)}", level=logging.ERROR) raise e @@ -136,10 +133,10 @@ This dataset appears clean with consistent formatting and no missing values, mak default_model_config = self._app_model_config else: default_model_config = { - "provider": os.getenv("MODEL_PROVIDER", "anthropic"), - "model": os.getenv("MODEL_NAME", "claude-sonnet-4-6"), - "api_key": os.getenv("ANTHROPIC_API_KEY"), - "temperature": _get_clamped_temperature(), + "provider": os.getenv("MODEL_PROVIDER", "openai"), + "model": os.getenv("MODEL_NAME", "gpt-4o-mini"), + "api_key": os.getenv("OPENAI_API_KEY"), + "temperature": float(os.getenv("TEMPERATURE", 1.0)), "max_tokens": int(os.getenv("MAX_TOKENS", 6000)) } @@ -147,21 +144,16 @@ This dataset appears clean with consistent formatting and no missing values, mak # Check if we need to create a brand new session logger.log_message(f"Creating new session state for session_id: {session_id}", level=logging.INFO) - # Initialize DuckDB connection for this session - - # Initialize with default state self._sessions[session_id] = { - "datasets": {"df":self._default_df.copy() if self._default_df is not None else None}, - "dataset_names": ["df"], + "current_df": self._default_df.copy() if self._default_df is not None else None, "retrievers": self._default_retrievers, "ai_system": self._default_ai_system, "make_data": self._make_data, "description": self._dataset_description, "name": self._default_name, "model_config": default_model_config, - "creation_time": time.time(), - "duckdb_conn": None, + "creation_time": time.time() } else: # Verify dataset integrity in existing session @@ -171,9 +163,9 @@ This dataset appears clean with consistent formatting and no missing values, mak session["model_config"] = default_model_config # If dataset is somehow missing, restore it - if "datasets" not in session or session["datasets"] is None: + if "current_df" not in session or session["current_df"] is None: logger.log_message(f"Restoring missing dataset for session {session_id}", level=logging.WARNING) - session["datasets"] = {"df":self._default_df.copy() if self._default_df is not None else None} + session["current_df"] = self._default_df.copy() if self._default_df is not None else None session["retrievers"] = self._default_retrievers session["ai_system"] = self._default_ai_system session["description"] = self._dataset_description @@ -190,61 +182,51 @@ This dataset appears clean with consistent formatting and no missing values, mak return self._sessions[session_id] - + def clear_session_state(self, session_id: str): + """ + Clear session-specific state + + Args: + session_id: The session identifier + """ + if session_id in self._sessions: + del self._sessions[session_id] - def update_session_dataset(self, session_id: str, datasets, names, desc: str, pre_generated=False): + def update_session_dataset(self, session_id: str, df, name: str, desc: str): """ - Update session with new dataset and optionally auto-generate description + Update dataset for a specific session + + Args: + session_id: The session identifier + df: Pandas DataFrame containing the dataset + name: Name of the dataset + desc: Description of the dataset """ try: + self._make_data = make_data(df, desc) + retrievers = self.initialize_retrievers(self.styling_instructions, [str(self._make_data)]) + ai_system = auto_analyst(agents=list(self.available_agents.values()), retrievers=retrievers) + # Get default model config for new sessions default_model_config = { - "provider": os.getenv("MODEL_PROVIDER", "anthropic"), - "model": os.getenv("MODEL_NAME", "claude-sonnet-4-6"), - "api_key": os.getenv("ANTHROPIC_API_KEY"), - "temperature": _get_clamped_temperature(), + "provider": os.getenv("MODEL_PROVIDER", "openai"), + "model": os.getenv("MODEL_NAME", "gpt-4o-mini"), + "api_key": os.getenv("OPENAI_API_KEY"), + "temperature": float(os.getenv("TEMPERATURE", 1.0)), "max_tokens": int(os.getenv("MAX_TOKENS", 6000)) } - # Get or create DuckDB connection for this session - - # Register the new dataset in DuckDB - - # Auto-generate description if we have datasets - if datasets and pre_generated==False: - try: - generated_desc = generate_dataset_description(datasets, desc, names) - desc = generated_desc # No need to format again since it's already formatted - logger.log_message(f"Auto-generated description for session {session_id}", level=logging.INFO) - except Exception as e: - logger.log_message(f"Failed to auto-generate description: {str(e)}", level=logging.WARNING) - # Keep the original description if generation fails - pass - - # Initialize retrievers and AI system BEFORE creating session_state - # Update make_data with the description - self._make_data = {'description': desc} - retrievers = self.initialize_retrievers(self.styling_instructions, [str(self._make_data)]) - - # Check if session has a user_id to create user-specific AI system - current_user_id = None - if session_id in self._sessions and "user_id" in self._sessions[session_id]: - current_user_id = self._sessions[session_id]["user_id"] - - ai_system = self.create_ai_system_for_user(retrievers, current_user_id) - # Create a completely fresh session state for the new dataset + # This ensures no remnants of the previous dataset remain session_state = { - "datasets": datasets, - "dataset_names": names, - "retrievers": retrievers, # Now retrievers is defined - "ai_system": ai_system, # Now ai_system is defined + "current_df": df, + "retrievers": retrievers, + "ai_system": ai_system, "make_data": self._make_data, "description": desc, - "name": names[0], - "duckdb_conn": None, - "model_config": default_model_config, + "name": name, + "model_config": default_model_config, # Initialize with default } # Preserve user_id, chat_id, and model_config if they exist in the current session @@ -254,12 +236,13 @@ This dataset appears clean with consistent formatting and no missing values, mak if "chat_id" in self._sessions[session_id]: session_state["chat_id"] = self._sessions[session_id]["chat_id"] if "model_config" in self._sessions[session_id]: + # Preserve the user's model configuration session_state["model_config"] = self._sessions[session_id]["model_config"] # Replace the entire session with the new state self._sessions[session_id] = session_state - logger.log_message(f"Updated session {session_id} with completely fresh dataset state: {str(names)}", level=logging.INFO) + logger.log_message(f"Updated session {session_id} with completely fresh dataset state: {name}", level=logging.INFO) except Exception as e: logger.log_message(f"Error updating dataset for session {session_id}: {str(e)}", level=logging.ERROR) raise e @@ -274,10 +257,10 @@ This dataset appears clean with consistent formatting and no missing values, mak try: # Get default model config from environment default_model_config = { - "provider": os.getenv("MODEL_PROVIDER", "anthropic"), - "model": os.getenv("MODEL_NAME", "claude-sonnet-4-6"), - "api_key": os.getenv("ANTHROPIC_API_KEY"), - "temperature": _get_clamped_temperature(), + "provider": os.getenv("MODEL_PROVIDER", "openai"), + "model": os.getenv("MODEL_NAME", "gpt-4o-mini"), + "api_key": os.getenv("OPENAI_API_KEY"), + "temperature": float(os.getenv("TEMPERATURE", 1.0)), "max_tokens": int(os.getenv("MAX_TOKENS", 6000)) } @@ -286,134 +269,21 @@ This dataset appears clean with consistent formatting and no missing values, mak del self._sessions[session_id] logger.log_message(f"Cleared existing state for session {session_id} before reset.", level=logging.INFO) - # Create new DuckDB connection for default session - # Initialize with default state self._sessions[session_id] = { - "datasets": {'df':self._default_df.copy()}, - "dataset_names": ["df"], # Use a copy + "current_df": self._default_df.copy(), # Use a copy "retrievers": self._default_retrievers, "ai_system": self._default_ai_system, "description": self._dataset_description, "name": self._default_name, # Explicitly set the default name "make_data": None, # Clear any custom make_data - "model_config": default_model_config, # Initialize with default model config - "duckdb_conn": None, # Create new DuckDB connection + "model_config": default_model_config # Initialize with default model config } logger.log_message(f"Reset session {session_id} to default dataset: {self._default_name}", level=logging.INFO) except Exception as e: logger.log_message(f"Error resetting session {session_id}: {str(e)}", level=logging.ERROR) raise e - def create_ai_system_for_user(self, retrievers, user_id=None): - """ - Create an AI system with user-specific agents (including custom agents) - - Args: - retrievers: The retrievers for the AI system - user_id: Optional user ID to load custom agents for - - Returns: - An auto_analyst instance with all available agents (standard + custom) - """ - try: - if user_id: - # Import here to avoid circular imports - from src.db.init_db import session_factory - - # Create a database session - db_session = session_factory() - try: - # Create AI system with user context to load custom agents - ai_system = auto_analyst( - agents=[], - retrievers=retrievers, - user_id=user_id, - db_session=db_session - ) - logger.log_message(f"Created AI system for user {user_id}", level=logging.INFO) - return ai_system - finally: - db_session.close() - else: - # Create standard AI system without custom agents - return auto_analyst(agents=[], retrievers=retrievers) - - except Exception as e: - logger.log_message(f"Error creating AI system for user {user_id}: {str(e)}", level=logging.ERROR) - # Fallback to standard AI system - return auto_analyst(agents=[], retrievers=retrievers) - - def set_default_lm_for_user(self, session_id: str, user_id: int = None): - """ - Set the default language model for a user upon signin using MODEL_OBJECTS. - - Args: - session_id: The session identifier - user_id: The authenticated user ID (optional) - - Returns: - Dictionary containing the default model configuration - """ - try: - # Import MODEL_OBJECTS directly - from src.utils.model_registry import MODEL_OBJECTS - - # Set Claude Sonnet 4.6 as default model - default_model_name = "claude-sonnet-4-6" - - # Ensure the model exists in MODEL_OBJECTS - if default_model_name not in MODEL_OBJECTS: - logger.log_message(f"Default model '{default_model_name}' not found in MODEL_OBJECTS, using gpt-5-mini", level=logging.WARNING) - default_model_name = "gpt-5-mini" - - # Get the model object directly from MODEL_OBJECTS - model_object = MODEL_OBJECTS[default_model_name] - - # Determine provider from model name - provider = "anthropic" # Claude models use Anthropic - - # Create default model configuration - default_model_config = { - "provider": provider, - "model": default_model_name, - "api_key": os.getenv(f"{provider.upper()}_API_KEY"), - "temperature": getattr(model_object, 'kwargs', {}).get('temperature', 0.7), - "max_tokens": getattr(model_object, 'kwargs', {}).get('max_tokens', 4000) - } - - # Ensure we have a session state for this session ID - if session_id not in self._sessions: - self.get_session_state(session_id) - - # Set the default model configuration in session state - self._sessions[session_id]["model_config"] = default_model_config - - # Also update the app-level model config if available - if hasattr(self, '_app_model_config'): - self._app_model_config.update(default_model_config) - - logger.log_message(f"Set default LM '{default_model_name}' for session {session_id} (user: {user_id})", level=logging.INFO) - - return { - "status": "success", - "model_config": default_model_config, - "message": f"Default model '{default_model_name}' set successfully" - } - - except Exception as e: - logger.log_message(f"Error setting default LM for user {user_id}: {str(e)}", level=logging.ERROR) - # Return fallback configuration - return { - "status": "error", - "model_config": { - "provider": "anthropic", - "model": "claude-sonnet-4-6", - "temperature": 0.7, - "max_tokens": 4000 - }, - "message": f"Failed to set default model, using fallback: {str(e)}" - } def set_session_user(self, session_id: str, user_id: int, chat_id: int = None): """ @@ -434,9 +304,6 @@ This dataset appears clean with consistent formatting and no missing values, mak # Store user ID self._sessions[session_id]["user_id"] = user_id - # Set default LM for user upon signin - self.set_default_lm_for_user(session_id, user_id) - # Generate or use chat ID if chat_id: chat_id_to_use = chat_id @@ -452,46 +319,33 @@ This dataset appears clean with consistent formatting and no missing values, mak # Store chat ID self._sessions[session_id]["chat_id"] = chat_id_to_use - # Recreate AI system with user context to load custom agents - try: - session_retrievers = self._sessions[session_id]["retrievers"] - user_ai_system = self.create_ai_system_for_user(session_retrievers, user_id) - self._sessions[session_id]["ai_system"] = user_ai_system - logger.log_message(f"Updated AI system for session {session_id} with user {user_id}", level=logging.INFO) - except Exception as e: - logger.log_message(f"Error updating AI system for user {user_id}: {str(e)}", level=logging.ERROR) - # Continue with existing AI system if update fails - # Make sure this data gets saved - logger.log_message(f"Associated session {session_id} with user {user_id}, chat_id: {chat_id_to_use}", level=logging.INFO) + logger.log_message(f"Associated session {session_id} with user_id={user_id}, chat_id={chat_id_to_use}", level=logging.INFO) # Return the updated session data return self._sessions[session_id] -async def get_session_id(request: Request, session_manager): - """ - Get or create a session ID from the request +async def get_session_id(request, session_manager): """ - # Debug: Log all headers - logger.log_message(f"🔍 ALL REQUEST HEADERS: {dict(request.headers)}", level=logging.DEBUG) + Get the session ID from the request, create/associate a user if needed - # Try to get session ID from headers FIRST (primary method) - session_id = request.headers.get("X-Session-ID") - logger.log_message(f"🔍 Session ID from X-Session-ID header: {session_id}", level=logging.DEBUG) + Args: + request: FastAPI Request object + session_manager: SessionManager instance + + Returns: + Session ID string + """ + # First try to get from query params + session_id = request.query_params.get("session_id") - # If not in headers, try query parameters (fallback for backward compatibility) + # If not in query params, try to get from headers if not session_id: - session_id = request.query_params.get("session_id") - logger.log_message(f"🔍 Session ID from query params: {session_id}", level=logging.DEBUG) - - logger.log_message(f"🔍 Final session_id before validation: '{session_id}' (type: {type(session_id)})", level=logging.DEBUG) + session_id = request.headers.get("X-Session-ID") - # STOP auto-generating sessions + # If still not found, generate a new one if not session_id: - logger.log_message(f"❌ No session ID found in request", level=logging.ERROR) - raise HTTPException(status_code=400, detail="Session ID required") - else: - logger.log_message(f"✅ Using existing session ID: {session_id}", level=logging.INFO) + session_id = str(uuid.uuid4()) # Get or create the session state session_state = session_manager.get_session_state(session_id) @@ -522,6 +376,24 @@ async def get_session_id(request: Request, session_manager): except (ValueError, TypeError): logger.log_message(f"Invalid user_id in query params: {user_id_param}", level=logging.WARNING) - # No user was found or created - just return the session ID without associating a user - logger.log_message(f"No authenticated user found for session {session_id}, continuing without user association", level=logging.INFO) + # Only create a guest user if no authenticated user is found + try: + # Create a guest user for this session + guest_username = f"guest_{session_id[:8]}" + guest_email = f"{guest_username}@example.com" + + # Create the user + user = create_user(username=guest_username, email=guest_email) + user_id = user.user_id + + logger.log_message(f"Created guest user {user_id} for session {session_id}", level=logging.INFO) + + # Associate the user with this session + session_manager.set_session_user( + session_id=session_id, + user_id=user_id + ) + except Exception as e: + logger.log_message(f"Error auto-creating user for session {session_id}: {str(e)}", level=logging.ERROR) + return session_id \ No newline at end of file diff --git a/src/managers/user_manager.py b/src/managers/user_manager.py index 3abd3b25a9257aafd4bd2c6bc5dcb76a26e4b148..75346339dae6f5bbcce3e8d7f468b53f03c99e05 100644 --- a/src/managers/user_manager.py +++ b/src/managers/user_manager.py @@ -1,14 +1,13 @@ import logging import os from typing import Optional -from datetime import datetime, UTC from fastapi import Depends, HTTPException, Request, status from fastapi.security import APIKeyHeader from src.db.init_db import get_session -from src.db.schemas.models import User as DBUser, AgentTemplate, UserTemplatePreference -from src.schemas.user_schema import User +from src.db.schemas.models import User as DBUser +from src.schemas.user_schemas import User from src.utils.logger import Logger logger = Logger("user_manager", see_time=True, console_log=False) @@ -26,16 +25,10 @@ async def get_current_user( Dependency to get the current authenticated user. Returns None if no user is authenticated. """ - # FastAPI resolves the `api_key` parameter when this function is used as a dependency. However, when the - # function is called directly (e.g. from the session manager), the `api_key` parameter will still hold the - # unresolved `Depends` placeholder object. In that case – or when no API key is supplied – we need to - # manually look for the key in the request headers or query parameters. - - if not api_key or not isinstance(api_key, str): - # Prefer header first for consistency with the dependency implementation - api_key = request.headers.get(API_KEY_NAME) or request.query_params.get("api_key") - - # If an API key still isn't available, treat the caller as anonymous + # If no API key is provided, return None (anonymous user) + if not api_key: + # Check for API key in query parameters (fallback) + api_key = request.query_params.get("api_key") if not api_key: return None @@ -101,9 +94,6 @@ def create_user(username: str, email: str) -> User: session.commit() session.refresh(new_user) - # Enable default agents for the new user - _enable_default_agents_for_user(new_user.user_id, session) - return User( user_id=new_user.user_id, username=new_user.username, @@ -123,8 +113,6 @@ def get_user_by_email(email: str) -> Optional[User]: session = get_session() try: user = session.query(DBUser).filter(DBUser.email == email).first() - if user is None: - return None return User( user_id=user.user_id, username=user.username, @@ -133,50 +121,3 @@ def get_user_by_email(email: str) -> Optional[User]: except Exception as e: logger.log_message(f"Error getting user by email: {str(e)}", logging.ERROR) return None - finally: - session.close() - -def _enable_default_agents_for_user(user_id: int, session): - """Enable default agents for a new user""" - try: - # Get all default agents (the 4 built-in agents) - default_agent_names = [ - "preprocessing_agent", - "statistical_analytics_agent", - "sk_learn_agent", - "data_viz_agent" - ] - - # Find these agents in the database - default_agents = session.query(AgentTemplate).filter( - AgentTemplate.template_name.in_(default_agent_names), - AgentTemplate.is_active == True - ).all() - - # Enable each default agent for the user - for agent in default_agents: - # Check if preference already exists - existing_pref = session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == agent.template_id - ).first() - - if not existing_pref: - # Create new preference with enabled=True - new_pref = UserTemplatePreference( - user_id=user_id, - template_id=agent.template_id, - is_enabled=True, # Enable by default - usage_count=0, - created_at=datetime.now(UTC), - updated_at=datetime.now(UTC) - ) - session.add(new_pref) - - session.commit() - logger.log_message(f"Enabled {len(default_agents)} default agents for user {user_id}", level=logging.INFO) - - except Exception as e: - session.rollback() - logger.log_message(f"Error enabling default agents for user {user_id}: {str(e)}", level=logging.ERROR) - raise diff --git a/src/routes/analytics_routes.py b/src/routes/analytics_routes.py index c23eafc66c5b21c9bd5ce492eceedf807d6b55e1..4f25b53d28db03f92b9f393f7558aef9521df5dd 100644 --- a/src/routes/analytics_routes.py +++ b/src/routes/analytics_routes.py @@ -2,7 +2,7 @@ import json import logging import os from collections import defaultdict -from datetime import datetime, timedelta, UTC +from datetime import datetime, timedelta from fastapi import APIRouter, Depends, HTTPException, Query, Request, WebSocket from fastapi.security import APIKeyHeader @@ -13,7 +13,7 @@ from sqlalchemy import case, desc, func from sqlalchemy.orm import Session from src.db.init_db import get_db, get_session -from src.db.schemas.models import ModelUsage, CodeExecution, Message, MessageFeedback, User +from src.db.schemas.models import ModelUsage from src.managers.chat_manager import ChatManager from typing import Any, Dict, List, Optional @@ -72,7 +72,7 @@ def get_model_tier(model_name): # Helper function to parse period parameter def get_date_range(period: str): - today = datetime.now(UTC) + today = datetime.utcnow() if period == '7d': start_date = today - timedelta(days=7) elif period == '30d': @@ -219,10 +219,6 @@ async def get_users( api_key: str = Depends(verify_admin_api_key) ): logger.log_message(f"User analytics requested with limit: {limit}, offset: {offset}", logging.INFO) - - # Only show users from the past 7 days - seven_days_ago = datetime.now(UTC) - timedelta(days=7) - user_query = db.query( ModelUsage.user_id, func.sum(ModelUsage.total_tokens).label("tokens"), @@ -231,8 +227,7 @@ async def get_users( func.min(ModelUsage.timestamp).label("first_seen"), func.max(ModelUsage.timestamp).label("last_seen") ).filter( - ModelUsage.user_id.isnot(None), - ModelUsage.timestamp >= seven_days_ago + ModelUsage.user_id.isnot(None) ).group_by( ModelUsage.user_id ).order_by( @@ -251,12 +246,9 @@ async def get_users( for user in user_query ] - # Get total users count for pagination (past 7 days only) + # Get total users count for pagination total_users = db.query(func.count(func.distinct(ModelUsage.user_id)))\ - .filter( - ModelUsage.user_id.isnot(None), - ModelUsage.timestamp >= seven_days_ago - )\ + .filter(ModelUsage.user_id.isnot(None))\ .scalar() or 0 logger.log_message(f"Retrieved {len(users)} users, total users: {total_users}", logging.INFO) @@ -351,7 +343,7 @@ async def get_session_stats( .scalar() or 0 # Active users today - today = datetime.now(UTC).date() + today = datetime.utcnow().date() active_today = db.query(func.count(func.distinct(ModelUsage.user_id)))\ .filter( func.date(ModelUsage.timestamp) == today, @@ -691,7 +683,7 @@ async def get_cost_projections( ): logger.log_message("Cost projections requested", logging.INFO) # Get last 30 days usage as baseline - thirty_days_ago = datetime.now(UTC) - timedelta(days=30) + thirty_days_ago = datetime.utcnow() - timedelta(days=30) baseline = db.query( func.sum(ModelUsage.cost).label("total_cost"), @@ -728,7 +720,7 @@ async def get_today_costs( api_key: str = Depends(verify_admin_api_key) ): logger.log_message("Today's costs requested", logging.INFO) - today = datetime.now(UTC).date() + today = datetime.utcnow().date() # Get today's costs today_data = db.query( @@ -1063,626 +1055,4 @@ async def get_tier_efficiency( "period": period, "start_date": tier_usage["start_date"], "end_date": tier_usage["end_date"] - } - -@router.get("/code-executions/summary") -async def get_code_execution_summary( - period: str = "30d", - db: Session = Depends(get_db), - api_key: str = Depends(verify_admin_api_key) -): - """ - Get a summary of code execution analytics including success rates, - common errors, and model performance. - """ - logger.log_message(f"Code execution summary requested for period: {period}", logging.INFO) - start_date, end_date = get_date_range(period) - - # Get overall code execution statistics - overall_stats = db.query( - func.count(CodeExecution.execution_id).label("total_executions"), - func.sum(case((CodeExecution.is_successful == True, 1), else_=0)).label("successful_executions"), - func.count(func.distinct(CodeExecution.user_id)).filter(CodeExecution.user_id.isnot(None)).label("total_users"), - func.count(func.distinct(CodeExecution.chat_id)).filter(CodeExecution.chat_id.isnot(None)).label("total_chats") - ).filter( - CodeExecution.created_at.between(start_date, end_date) - ).first() - - # Calculate success rate - total_executions = overall_stats.total_executions or 0 - successful_executions = overall_stats.successful_executions or 0 - success_rate = (successful_executions / total_executions) if total_executions > 0 else 0 - - # Get model performance stats - model_stats = db.query( - CodeExecution.model_name, - func.count(CodeExecution.execution_id).label("total_executions"), - func.sum(case((CodeExecution.is_successful == True, 1), else_=0)).label("successful_executions"), - func.count(func.distinct(CodeExecution.user_id)).filter(CodeExecution.user_id.isnot(None)).label("user_count") - ).filter( - CodeExecution.created_at.between(start_date, end_date), - CodeExecution.model_name.isnot(None) - ).group_by( - CodeExecution.model_name - ).order_by( - desc(func.count(CodeExecution.execution_id)) - ).all() - - # Process model stats - model_performance = [] - for model in model_stats: - model_executions = model.total_executions or 0 - model_success_rate = (model.successful_executions / model.total_executions) if model.total_executions > 0 else 0 - model_performance.append({ - "model_name": model.model_name, - "total_executions": model_executions, - "successful_executions": model.successful_executions or 0, - "success_rate": model_success_rate, - "user_count": model.user_count or 0 - }) - - # Get failed agent statistics - failed_agent_data = [] - failed_executions = db.query(CodeExecution).filter( - CodeExecution.created_at.between(start_date, end_date), - CodeExecution.is_successful == False, - CodeExecution.failed_agents.isnot(None) - ).all() - - # Count agent failures - agent_failure_counts = {} - for execution in failed_executions: - try: - if execution.failed_agents: - failed_agents = json.loads(execution.failed_agents) - for agent in failed_agents: - agent_failure_counts[agent] = agent_failure_counts.get(agent, 0) + 1 - except (json.JSONDecodeError, TypeError): - logger.log_message(f"Error parsing failed_agents JSON: {execution.failed_agents}", logging.ERROR) - - # Convert to list for response - failed_agent_data = [ - {"agent_name": agent, "failure_count": count} - for agent, count in sorted(agent_failure_counts.items(), key=lambda x: x[1], reverse=True) - ] - - # Return the compiled data - result = { - "period": period, - "start_date": start_date.strftime('%Y-%m-%d'), - "end_date": end_date.strftime('%Y-%m-%d'), - "overall_stats": { - "total_executions": total_executions, - "successful_executions": successful_executions, - "failed_executions": total_executions - successful_executions, - "success_rate": success_rate, - "total_users": overall_stats.total_users or 0, - "total_chats": overall_stats.total_chats or 0 - }, - "model_performance": model_performance, - "failed_agents": failed_agent_data - } - - logger.log_message(f"Code execution summary retrieved with {total_executions} executions", logging.INFO) - return result - -@router.get("/code-executions/detailed") -async def get_detailed_code_executions( - period: str = "30d", - success_filter: Optional[bool] = None, - user_id: Optional[int] = None, - model_name: Optional[str] = None, - limit: int = 100, - db: Session = Depends(get_db), - api_key: str = Depends(verify_admin_api_key) -): - """ - Get detailed information about individual code executions - with filtering options. - """ - logger.log_message(f"Detailed code executions requested for period: {period}", logging.INFO) - - try: - start_date, end_date = get_date_range(period) - - # Add this at the beginning of the function, right after the logger.log_message line: - logger.log_message(f"Database session type: {type(db)}", logging.INFO) - logger.log_message(f"Database URL: {db.bind.url if hasattr(db, 'bind') and db.bind else 'No bind'}", logging.INFO) - - # Test if we can access the CodeExecution model - try: - test_count = db.query(CodeExecution).count() - logger.log_message(f"Total CodeExecution records in database: {test_count}", logging.INFO) - except Exception as test_error: - logger.log_message(f"Database connection test failed: {str(test_error)}", logging.ERROR) - raise HTTPException( - status_code=500, - detail=f"Database connection test failed: {str(test_error)}" - ) - - # Build the query with filters - query = db.query(CodeExecution).filter( - CodeExecution.created_at.between(start_date, end_date) - ) - - # Apply optional filters - if success_filter is not None: - query = query.filter(CodeExecution.is_successful == success_filter) - - if user_id: - query = query.filter(CodeExecution.user_id == user_id) - - if model_name: - query = query.filter(CodeExecution.model_name == model_name) - - # Order by most recent first and limit results - query = query.order_by(desc(CodeExecution.created_at)).limit(limit) - - # Execute query with error handling - try: - executions = query.all() - except Exception as db_error: - logger.log_message(f"Database query error: {str(db_error)}", logging.ERROR) - raise HTTPException( - status_code=500, - detail=f"Database query failed: {str(db_error)}" - ) - - # Process results with defensive programming - detailed_executions = [] - for execution in executions: - # Skip None or invalid execution objects - if not execution or not hasattr(execution, 'execution_id'): - logger.log_message(f"Skipping invalid execution object: {execution}", logging.WARNING) - continue - - # Parse failed agents and error messages if available - failed_agents_list = [] - error_messages_dict = {} - - try: - if execution.failed_agents: - failed_agents_list = json.loads(execution.failed_agents) - - if execution.error_messages: - error_messages_dict = json.loads(execution.error_messages) - except (json.JSONDecodeError, TypeError): - logger.log_message(f"Error parsing JSON data for execution {execution.execution_id}", logging.ERROR) - - # Format the execution data with null checks - try: - detailed_executions.append({ - "execution_id": execution.execution_id, - "message_id": execution.message_id, - "chat_id": execution.chat_id, - "user_id": execution.user_id, - "created_at": execution.created_at.isoformat() if execution.created_at else None, - "updated_at": execution.updated_at.isoformat() if execution.updated_at else None, - "is_successful": getattr(execution, 'is_successful', False), # Safe attribute access - "model_info": { - "provider": getattr(execution, 'model_provider', None), - "name": getattr(execution, 'model_name', None), - "temperature": getattr(execution, 'model_temperature', None), - "max_tokens": getattr(execution, 'model_max_tokens', None) - }, - "failed_agents": failed_agents_list, - "error_messages": error_messages_dict, - # Include trimmed code snippets with null checks - "initial_code_preview": execution.initial_code[:500] + "..." if execution.initial_code and len(execution.initial_code) > 500 else execution.initial_code, - "latest_code_preview": execution.latest_code[:500] + "..." if execution.latest_code and len(execution.latest_code) > 500 else execution.latest_code, - }) - except Exception as e: - logger.log_message(f"Error processing execution {execution.execution_id}: {str(e)}", logging.ERROR) - continue - - logger.log_message(f"Retrieved {len(detailed_executions)} detailed code executions", logging.INFO) - return { - "period": period, - "start_date": start_date.strftime('%Y-%m-%d'), - "end_date": end_date.strftime('%Y-%m-%d'), - "count": len(detailed_executions), - "executions": detailed_executions - } - - except Exception as e: - logger.log_message(f"Database error in get_detailed_code_executions: {str(e)}", logging.ERROR) - raise HTTPException( - status_code=500, - detail=f"Database error: {str(e)}" - ) - -@router.get("/code-executions/users") -async def get_user_code_execution_stats( - period: str = "30d", - limit: int = 50, - db: Session = Depends(get_db), - api_key: str = Depends(verify_admin_api_key) -): - """ - Get statistics about code executions grouped by user - """ - logger.log_message(f"User code execution stats requested for period: {period}", logging.INFO) - start_date, end_date = get_date_range(period) - - # Get user statistics for code executions - user_stats = db.query( - CodeExecution.user_id, - func.count(CodeExecution.execution_id).label("total_executions"), - func.sum(case((CodeExecution.is_successful == True, 1), else_=0)).label("successful_executions"), - func.min(CodeExecution.created_at).label("first_execution"), - func.max(CodeExecution.created_at).label("latest_execution") - ).filter( - CodeExecution.created_at.between(start_date, end_date), - CodeExecution.user_id.isnot(None) - ).group_by( - CodeExecution.user_id - ).order_by( - desc(func.count(CodeExecution.execution_id)) - ).limit(limit).all() - - # Process user statistics - users_data = [] - for user in user_stats: - success_rate = (user.successful_executions / user.total_executions) if user.total_executions > 0 else 0 - users_data.append({ - "user_id": user.user_id, - "total_executions": user.total_executions, - "successful_executions": user.successful_executions, - "failed_executions": user.total_executions - user.successful_executions, - "success_rate": success_rate, - "first_execution": user.first_execution.isoformat() if user.first_execution else None, - "latest_execution": user.latest_execution.isoformat() if user.latest_execution else None - }) - - logger.log_message(f"Retrieved code execution stats for {len(users_data)} users", logging.INFO) - return { - "period": period, - "start_date": start_date.strftime('%Y-%m-%d'), - "end_date": end_date.strftime('%Y-%m-%d'), - "users": users_data - } - -@router.get("/code-executions/error-analysis") -async def get_error_analysis( - period: str = "30d", - db: Session = Depends(get_db), - api_key: str = Depends(verify_admin_api_key) -): - """ - Analyze common error patterns in failed code executions - """ - logger.log_message(f"Code execution error analysis requested for period: {period}", logging.INFO) - start_date, end_date = get_date_range(period) - - # Get failed executions - failed_executions = db.query(CodeExecution).filter( - CodeExecution.created_at.between(start_date, end_date), - CodeExecution.is_successful == False, - CodeExecution.error_messages.isnot(None) - ).all() - - # Analyze error messages and categorize them - error_types = {} - error_by_agent = {} - - for execution in failed_executions: - try: - if execution.error_messages: - error_messages = json.loads(execution.error_messages) - for agent, error in error_messages.items(): - # Add to agent-specific counts - if agent not in error_by_agent: - error_by_agent[agent] = { - "count": 0, - "common_errors": {} - } - error_by_agent[agent]["count"] += 1 - - # Categorize the error - error_category = categorize_error(error) - error_by_agent[agent]["common_errors"][error_category] = error_by_agent[agent]["common_errors"].get(error_category, 0) + 1 - - # Add to overall error type counts - error_types[error_category] = error_types.get(error_category, 0) + 1 - except (json.JSONDecodeError, TypeError): - logger.log_message(f"Error parsing error_messages JSON: {execution.error_messages}", logging.ERROR) - - # Convert to lists for response - error_types_list = [ - {"error_type": error_type, "count": count} - for error_type, count in sorted(error_types.items(), key=lambda x: x[1], reverse=True) - ] - - error_by_agent_list = [ - { - "agent_name": agent, - "error_count": data["count"], - "common_errors": [ - {"error_type": error_type, "count": count} - for error_type, count in sorted(data["common_errors"].items(), key=lambda x: x[1], reverse=True) - ] - } - for agent, data in sorted(error_by_agent.items(), key=lambda x: x[1]["count"], reverse=True) - ] - - logger.log_message(f"Analyzed errors from {len(failed_executions)} failed executions", logging.INFO) - return { - "period": period, - "start_date": start_date.strftime('%Y-%m-%d'), - "end_date": end_date.strftime('%Y-%m-%d'), - "total_failed_executions": len(failed_executions), - "error_types": error_types_list, - "error_by_agent": error_by_agent_list - } - -# Helper function to categorize error messages -def categorize_error(error_message): - """ - Categorize an error message into common types - """ - error_message = str(error_message).lower() - - if "nameerror" in error_message or "name '" in error_message: - return "NameError" - elif "syntaxerror" in error_message: - return "SyntaxError" - elif "typeerror" in error_message: - return "TypeError" - elif "attributeerror" in error_message: - return "AttributeError" - elif "indexerror" in error_message or "keyerror" in error_message: - return "IndexError/KeyError" - elif "importerror" in error_message or "modulenotfounderror" in error_message: - return "ImportError" - elif "valueerror" in error_message: - return "ValueError" - elif "unsupported operand" in error_message: - return "OperationError" - elif "indent" in error_message: - return "IndentationError" - elif "permission" in error_message: - return "PermissionError" - elif "filenotfound" in error_message: - return "FileNotFoundError" - elif "memory" in error_message: - return "MemoryError" - elif "timeout" in error_message or "timed out" in error_message: - return "TimeoutError" - else: - return "OtherError" - -@router.get("/feedback/summary") -async def get_feedback_summary( - period: str = "30d", - db: Session = Depends(get_db), - api_key: str = Depends(verify_admin_api_key) -): - """ - Get summary statistics about user feedback ratings. - """ - logger.log_message(f"Feedback summary requested for period: {period}", logging.INFO) - start_date, end_date = get_date_range(period) - - # Get overall feedback statistics - overall_stats = db.query( - func.count().label("total_feedback"), - func.avg(MessageFeedback.rating).label("avg_rating"), - func.count(func.distinct(Message.chat_id)).label("chats_with_feedback") - ).join( - Message, MessageFeedback.message_id == Message.message_id - ).filter( - MessageFeedback.created_at.between(start_date, end_date), - MessageFeedback.rating.isnot(None) - ).first() - - # Get feedback counts by rating - rating_counts = db.query( - MessageFeedback.rating, - func.count().label("count") - ).filter( - MessageFeedback.created_at.between(start_date, end_date), - MessageFeedback.rating.isnot(None) - ).group_by( - MessageFeedback.rating - ).all() - - # Convert to dictionary for easier access - ratings_dict = {rating.rating: rating.count for rating in rating_counts} - - # Ensure all ratings (1-5) are represented - ratings_distribution = [ - {"rating": i, "count": ratings_dict.get(i, 0)} - for i in range(1, 6) - ] - - # Get feedback by model - model_feedback = db.query( - MessageFeedback.model_name, - func.count().label("total_feedback"), - func.avg(MessageFeedback.rating).label("avg_rating") - ).filter( - MessageFeedback.created_at.between(start_date, end_date), - MessageFeedback.rating.isnot(None), - MessageFeedback.model_name.isnot(None) - ).group_by( - MessageFeedback.model_name - ).order_by( - desc(func.avg(MessageFeedback.rating)) - ).all() - - # Process model feedback - models_data = [ - { - "model_name": model.model_name, - "total_feedback": model.total_feedback, - "avg_rating": float(model.avg_rating) if model.avg_rating else 0 - } - for model in model_feedback - ] - - # Get feedback trend over time - daily_feedback = db.query( - func.date(MessageFeedback.created_at).label("date"), - func.avg(MessageFeedback.rating).label("avg_rating"), - func.count().label("count") - ).filter( - MessageFeedback.created_at.between(start_date, end_date), - MessageFeedback.rating.isnot(None) - ).group_by( - func.date(MessageFeedback.created_at) - ).order_by( - func.date(MessageFeedback.created_at) - ).all() - - # Process daily feedback - feedback_trend = [ - { - "date": str(day.date), - "avg_rating": float(day.avg_rating) if day.avg_rating else 0, - "count": day.count - } - for day in daily_feedback - ] - - # Fill in any missing dates with null values - date_range = [(start_date + timedelta(days=i)).strftime('%Y-%m-%d') - for i in range((end_date - start_date).days + 1)] - - feedback_dict = {item["date"]: item for item in feedback_trend} - - filled_trend = [] - for date in date_range: - if date in feedback_dict: - filled_trend.append(feedback_dict[date]) - else: - filled_trend.append({ - "date": date, - "avg_rating": None, - "count": 0 - }) - - logger.log_message(f"Feedback summary retrieved with {overall_stats.total_feedback or 0} ratings", logging.INFO) - return { - "period": period, - "start_date": start_date.strftime('%Y-%m-%d'), - "end_date": end_date.strftime('%Y-%m-%d'), - "total_feedback": overall_stats.total_feedback or 0, - "avg_rating": float(overall_stats.avg_rating) if overall_stats.avg_rating else 0, - "chats_with_feedback": overall_stats.chats_with_feedback or 0, - "ratings_distribution": ratings_distribution, - "models_data": models_data, - "feedback_trend": filled_trend - } - -@router.get("/feedback/detailed") -async def get_detailed_feedback( - period: str = "30d", - min_rating: Optional[int] = None, - max_rating: Optional[int] = None, - model_name: Optional[str] = None, - limit: int = 100, - offset: int = 0, - db: Session = Depends(get_db), - api_key: str = Depends(verify_admin_api_key) -): - """ - Get detailed feedback records with filtering options - """ - logger.log_message(f"Detailed feedback requested for period: {period}", logging.INFO) - start_date, end_date = get_date_range(period) - - # Build query with joins to get message content - query = db.query( - MessageFeedback, - Message.content.label("message_content"), - Message.sender.label("message_sender"), - Message.chat_id.label("chat_id") - ).join( - Message, MessageFeedback.message_id == Message.message_id - ).filter( - MessageFeedback.created_at.between(start_date, end_date), - MessageFeedback.rating.isnot(None) - ) - - # Apply optional filters - if min_rating is not None: - query = query.filter(MessageFeedback.rating >= min_rating) - - if max_rating is not None: - query = query.filter(MessageFeedback.rating <= max_rating) - - if model_name: - query = query.filter(MessageFeedback.model_name == model_name) - - # Get count for pagination before adding limit/offset - total_count = query.count() - - # Apply pagination - query = query.order_by(desc(MessageFeedback.created_at)).offset(offset).limit(limit) - - # Execute query - results = query.all() - - # Process results - detailed_feedback = [] - for result in results: - feedback = result.MessageFeedback - - detailed_feedback.append({ - "feedback_id": feedback.feedback_id, - "message_id": feedback.message_id, - "chat_id": result.chat_id, - "rating": feedback.rating, - "model_name": feedback.model_name, - "model_provider": feedback.model_provider, - "temperature": feedback.temperature, - "max_tokens": feedback.max_tokens, - "created_at": feedback.created_at.isoformat() if feedback.created_at else None, - "message_preview": result.message_content[:200] + "..." if len(result.message_content) > 200 else result.message_content, - "message_sender": result.message_sender - }) - - logger.log_message(f"Retrieved {len(detailed_feedback)} detailed feedback records", logging.INFO) - return { - "period": period, - "start_date": start_date.strftime('%Y-%m-%d'), - "end_date": end_date.strftime('%Y-%m-%d'), - "total": total_count, - "count": len(detailed_feedback), - "offset": offset, - "limit": limit, - "feedback": detailed_feedback - } - -@router.get("/public/ticker") -async def get_public_ticker_data(db: Session = Depends(get_db)): - """ - Get public ticker data for the landing page showing overall platform statistics. - This endpoint is public and doesn't require authentication. - """ - try: - # Get total number of users (signups) - total_signups = db.query(func.count(func.distinct(User.user_id))).scalar() or 0 - - # Get total tokens used across all model usage - total_tokens = db.query(func.sum(ModelUsage.total_tokens)).scalar() or 0 - - # Get total requests made - total_requests = db.query(func.count(ModelUsage.usage_id)).scalar() or 0 - - # Return the ticker data - return { - "total_signups": total_signups, - "total_tokens": int(total_tokens), - "total_requests": total_requests, - "last_updated": datetime.now(UTC).isoformat() - } - except Exception as e: - logger.log_message(f"Error retrieving public ticker data: {str(e)}", logging.ERROR) - # Return default values in case of error - return { - "total_signups": 0, - "total_tokens": 0, - "total_requests": 0, - "last_updated": datetime.now(UTC).isoformat() - } \ No newline at end of file + } \ No newline at end of file diff --git a/src/routes/blog_routes.py b/src/routes/blog_routes.py deleted file mode 100644 index be27b5f2f001046e6ef705db8d68cf6f6bb598fd..0000000000000000000000000000000000000000 --- a/src/routes/blog_routes.py +++ /dev/null @@ -1,96 +0,0 @@ -import os -import json -from fastapi import APIRouter, HTTPException -from typing import List -from pydantic import BaseModel - -router = APIRouter() - -class BlogPost(BaseModel): - id: str - title: str - excerpt: str - content: str - author: str - publishedAt: str - tags: List[str] - featured: bool - readTime: str - -@router.get("/api/blog/posts", response_model=List[BlogPost]) -async def get_blog_posts(): - """Get all blog posts""" - try: - # Get the path to the utils/data directory - current_dir = os.path.dirname(os.path.abspath(__file__)) - utils_dir = os.path.join(current_dir, '..', 'utils', 'data') - json_path = os.path.join(utils_dir, 'sample-posts.json') - - # Normalize the path - json_path = os.path.normpath(json_path) - - if not os.path.exists(json_path): - raise HTTPException(status_code=404, detail=f"Blog posts data file not found at: {json_path}") - - with open(json_path, 'r', encoding='utf-8') as f: - posts = json.load(f) - - return posts - - except FileNotFoundError: - raise HTTPException(status_code=404, detail=f"Blog posts data file not found at: {json_path}") - except json.JSONDecodeError: - raise HTTPException(status_code=500, detail="Invalid JSON format in blog posts data") - except Exception as e: - raise HTTPException(status_code=500, detail=f"Error loading blog posts: {str(e)}") - -@router.get("/api/blog/posts/{post_id}", response_model=BlogPost) -async def get_blog_post(post_id: str): - """Get a specific blog post by ID""" - try: - posts = await get_blog_posts() - - for post in posts: - if post['id'] == post_id: - return post - - raise HTTPException(status_code=404, detail="Blog post not found") - - except HTTPException: - raise - except Exception as e: - raise HTTPException(status_code=500, detail=f"Error loading blog post: {str(e)}") - -@router.get("/api/blog/posts/featured", response_model=BlogPost) -async def get_featured_post(): - """Get the featured blog post""" - try: - posts = await get_blog_posts() - - for post in posts: - if post.get('featured', False): - return post - - raise HTTPException(status_code=404, detail="No featured post found") - - except HTTPException: - raise - except Exception as e: - raise HTTPException(status_code=500, detail=f"Error loading featured post: {str(e)}") - -@router.get("/api/blog/tags") -async def get_blog_tags(): - """Get all unique tags from blog posts""" - try: - posts = await get_blog_posts() - - all_tags = set() - for post in posts: - all_tags.update(post.get('tags', [])) - - return list(all_tags) - - except HTTPException: - raise - except Exception as e: - raise HTTPException(status_code=500, detail=f"Error loading blog tags: {str(e)}") diff --git a/src/routes/chat_routes.py b/src/routes/chat_routes.py index 80f378d0d86acb89145a487f714b0b3212874fef..d44c487e37271fbc6c70cd7c1bceac83f013f565 100644 --- a/src/routes/chat_routes.py +++ b/src/routes/chat_routes.py @@ -8,7 +8,7 @@ from src.db.schemas.models import ModelUsage from src.managers.ai_manager import AI_Manager from src.managers.chat_manager import ChatManager from src.managers.user_manager import get_current_user, User -from src.schemas.chat_schema import * +from src.schemas.chat_schemas import * from src.utils.logger import Logger import os from dotenv import load_dotenv @@ -132,3 +132,49 @@ async def cleanup_empty_chats(request: ChatCreate): except Exception as e: logger.log_message(f"Error cleaning up empty chats: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=500, detail=f"Failed to clean up empty chats: {str(e)}") + +@router.post("/debug/test-model-usage") +async def test_model_usage( + model_name: str = "gpt-3.5-turbo", + user_id: Optional[int] = None +): + """Debug endpoint to manually test model usage tracking""" + try: + # Generate a test prompt + test_prompt = "This is a test message to verify model usage tracking." + + # Call the AI manager directly + response = await ai_manager.generate_response( + prompt=test_prompt, + model_name=model_name, + user_id=user_id, + chat_id=999 # Test chat ID + ) + + # Get the latest model usage entry + session = session_factory() + try: + latest_usage = session.query(ModelUsage).order_by(ModelUsage.usage_id.desc()).first() + + return { + "success": True, + "message": "Model usage tracking test completed", + "response": response, + "usage_recorded": { + "usage_id": latest_usage.usage_id if latest_usage else None, + "model_name": latest_usage.model_name if latest_usage else None, + "tokens": latest_usage.total_tokens if latest_usage else None, + "cost": latest_usage.cost if latest_usage else None, + "timestamp": latest_usage.timestamp.isoformat() if latest_usage else None + } + } + finally: + session.close() + + except Exception as e: + logger.log_message(f"Error in test-model-usage: {str(e)}", level=logging.ERROR) + return { + "success": False, + "error": str(e) + } + \ No newline at end of file diff --git a/src/routes/code_routes.py b/src/routes/code_routes.py index 660971e806deef77bc537795fdf3ad0834dc8c21..9c8b7b352e3689a07934cd2117e91f35cc324e89 100644 --- a/src/routes/code_routes.py +++ b/src/routes/code_routes.py @@ -1,7 +1,6 @@ import io import logging import re -import json from fastapi import APIRouter, Depends, HTTPException, Request from typing import Dict, Optional, List, Tuple from pydantic import BaseModel @@ -10,47 +9,8 @@ from scripts.format_response import execute_code_from_markdown, format_code_bloc from src.utils.logger import Logger from src.routes.session_routes import get_session_id_dependency from src.agents.agents import code_edit, code_fix -from src.db.schemas.models import CodeExecution -from src.db.init_db import get_session import dspy -import textwrap import os -from src.schemas.code_schema import CodeExecuteRequest, CodeEditRequest, CodeFixRequest, CodeCleanRequest, GetLatestCodeRequest -from src.utils.model_registry import MODEL_OBJECTS -import asyncio -import traceback - -def clean_print_statements(code_block): - """ - This function cleans up any `print()` statements that might contain unwanted `\n` characters. - It ensures print statements are properly formatted without unnecessary newlines. - """ - # This regex targets print statements, even if they have newlines inside - return re.sub(r'print\((.*?)(\\n.*?)(.*?)\)', r'print(\1\3)', code_block, flags=re.DOTALL) - - -def remove_main_block(code): - # Match the __main__ block - pattern = r'(?m)^if\s+__name__\s*==\s*["\']__main__["\']\s*:\s*\n((?:\s+.*\n?)*)' - - match = re.search(pattern, code) - if match: - main_block = match.group(1) - - # Dedent the code block inside __main__ - dedented_block = textwrap.dedent(main_block) - - # Remove \n from any print statements in the block (also handling multiline print cases) - dedented_block = clean_print_statements(dedented_block) - # Replace the block in the code - cleaned_code = re.sub(pattern, dedented_block, code) - - # Optional: Remove leading newlines if any - cleaned_code = cleaned_code.strip() - - return cleaned_code - return code - # Initialize router router = APIRouter( prefix="/code", @@ -61,45 +21,21 @@ router = APIRouter( # Initialize logger logger = Logger("code_routes", see_time=True, console_log=False) try_logger = Logger("try_code_routes", see_time=True, console_log=False) - - -def score_code(args, code, datasets=None): - """ - Simple code scorer that checks if code runs successfully. - Handles both deep analysis (combined_code) and fix (fixed_code) scenarios. +# Request body model +class CodeExecuteRequest(BaseModel): + code: str - Args: - args: Arguments (unused but required for dspy.Refine) - code: Code object with combined_code, fixed_code, or string content - datasets: Dictionary of datasets from session state (optional) +class CodeEditRequest(BaseModel): + original_code: str + user_prompt: str + +class CodeFixRequest(BaseModel): + code: str + error: str + +class CodeCleanRequest(BaseModel): + code: str - Returns: - int: Score (0=error, 1=success) - """ - try: - # Handle different attribute names based on context - if hasattr(code, 'combined_code'): - code_to_execute = code.combined_code # Deep analysis - elif hasattr(code, 'fixed_code'): - code_to_execute = code.fixed_code # Code fix - else: - code_to_execute = str(code) # Fallback to string - - # Make datasets available if provided - local_vars = {} - if datasets: - local_vars.update(datasets) - - # Execute the code to test if it works - exec(code_to_execute, local_vars) - return 1 # Success - except Exception as e: - return 0 # Error - - -# Remove the global refine_fixer declaration - - def format_code(code: str) -> str: """ Clean the code by organizing imports and ensuring code blocks are properly formatted. @@ -185,14 +121,8 @@ def identify_error_blocks(code: str, error_output: str) -> List[Tuple[str, str, # Parse the error output to find which agents had errors faulty_blocks = [] - # Find error patterns like "=== ERROR IN AGENT_NAME ===" or "=== ERROR IN UNKNOWN_AGENT ===" - error_matches = [] - for match in re.finditer( - r'^===\s+ERROR\s+IN\s+([A-Za-z0-9_]+)\s+===\s*([\s\S]*?)(?=^===\s+[A-Z]+\s+IN\s+[A-Za-z0-9_]+\s+===|\Z)', - error_output, - re.MULTILINE - ): - error_matches.append((match.group(1), match.group(2))) + # Find error patterns like "=== ERROR IN AGENT_NAME ===" + error_matches = re.findall(r'===\s+ERROR\s+IN\s+([A-Za-z0-9_]+)\s+===\s*([\s\S]*?)(?:(?===\s+)|$)', error_output) if not error_matches: return [] @@ -269,68 +199,113 @@ def extract_relevant_error_section(error_message: str) -> str: # If we couldn't find "Problem at this location", include first few and last few lines if len(error_lines) > 10: - return '\n'.join(error_lines[:5] + error_lines[-7:]) + return '\n'.join(error_lines[:3] + error_lines[-3:]) # If the error is short enough, return as is return error_message -async def fix_code_with_dspy(code: str, error: str, dataset_context: str = "", datasets: dict = None): +def fix_code_with_dspy(code: str, error: str, dataset_context: str = ""): """ - Fix code using DSPy Refine with datasets-aware reward function - """ - try: - # Wrap score_code to fix datasets argument - reward_fn_with_datasets = lambda args, pred: score_code(args, pred, datasets=datasets) - - refine_fixer = dspy.Refine( - module=dspy.Predict(code_fix), - N=3, - threshold=1.0, - reward_fn=reward_fn_with_datasets, - fail_count=3 - ) - - # Check if we have valid API key - anthropic_key = os.environ.get('ANTHROPIC_API_KEY') - if not anthropic_key: - raise ValueError("ANTHROPIC_API_KEY environment variable is not set") + Fix code with errors by identifying faulty blocks and fixing them individually + + Args: + code (str): The code containing errors + error (str): Error message from execution + dataset_context (str): Context about the dataset - # Fix the entire code using refine - try: - # Create the LM instance that will be used - thread_lm = MODEL_OBJECTS['claude-sonnet-4-6'] - - # Define the blocking function to run in thread - def run_refine_fixer(): - with dspy.context(lm=thread_lm): - return refine_fixer( - dataset_context=str(dataset_context) or "", - faulty_code=str(code) or "", - error=str(error) or "", - ) - - # Use asyncio.to_thread for better async integration with timeout - result = await asyncio.wait_for( - asyncio.to_thread(run_refine_fixer), - timeout=60.0 # 60 second timeout + Returns: + str: The fixed code + """ + gemini = dspy.LM("gemini/gemini-2.5-pro-preview-03-25", api_key = os.environ['GEMINI_API_KEY'], max_tokens=5000) + + # Find the blocks with errors + faulty_blocks = identify_error_blocks(code, error) + logger.log_message(f"Number of faulty blocks found: {len(faulty_blocks)}", level=logging.INFO) + if not faulty_blocks: + # If no specific errors found, fix the entire code + with dspy.context(lm=gemini): + code_fixer = dspy.ChainOfThought(code_fix) + result = code_fixer( + dataset_context=str(dataset_context) or "", + faulty_code=str(code) or "", + error=str(error) or "", ) - - if not hasattr(result, 'fixed_code'): - raise ValueError("DSPy Refine did not return a result with 'fixed_code' attribute") - return result.fixed_code + + # Start with the original code + result_code = code.replace("```python", "").replace("```", "") + + # Fix each faulty block separatelyw + with dspy.context(lm=gemini): + code_fixer = dspy.ChainOfThought(code_fix) + + for agent_name, block_code, specific_error in faulty_blocks: + logger.log_message(f"Fixing {agent_name} block", level=logging.INFO) - except Exception as e: - logger.log_message(f"🔧 DETAILED ERROR in fix_code_with_dspy: {str(e)}", level=logging.ERROR) - logger.log_message(f"�� ERROR TYPE: {type(e).__name__}", level=logging.ERROR) - logger.log_message(f"�� ERROR TRACEBACK: {traceback.format_exc()}", level=logging.ERROR) - - # Instead of returning original code, raise the error so we can see what's wrong - raise HTTPException(status_code=500, detail=f"Fix failed: {str(e)}") - - except Exception as e: - logger.log_message(f"Error in fix_code_with_dspy: {str(e)}", level=logging.ERROR) - raise RuntimeError(f"Fix code setup failed: {str(e)}") from e + try: + # Extract inner code between the markers + inner_code_match = re.search(r'#\s+\w+\s+code\s+start\s*\n([\s\S]*?)#\s+\w+\s+code\s+end', block_code) + if not inner_code_match: + logger.log_message(f"Could not extract inner code for {agent_name}", level=logging.WARNING) + continue + + inner_code = inner_code_match.group(1).strip() + + # Find markers + start_marker_match = re.search(r'(#\s+\w+\s+code\s+start)', block_code) + end_marker_match = re.search(r'(#\s+\w+\s+code\s+end)', block_code) + + if not start_marker_match or not end_marker_match: + logger.log_message(f"Could not find start/end markers for {agent_name}", level=logging.WARNING) + continue + + start_marker = start_marker_match.group(1) + end_marker = end_marker_match.group(1) + + # Extract the error type and actual error message + error_type = "" + error_msg = specific_error + + # Look for common error patterns to provide focused context to the LLM + error_type_match = re.search(r'(TypeError|ValueError|AttributeError|IndexError|KeyError|NameError):\s*([^\n]+)', specific_error) + if error_type_match: + error_type = error_type_match.group(1) + error_msg = f"{error_type}: {error_type_match.group(2)}" + + # Add problem location if available + if "Problem at this location:" in specific_error: + problem_section = re.search(r'Problem at this location:([\s\S]*?)(?:\n\n|$)', specific_error) + if problem_section: + error_msg = f"{error_msg}\n\nProblem at: {problem_section.group(1).strip()}" + + # Fix only the inner code + result = code_fixer( + dataset_context=str(dataset_context) or "", + faulty_code=str(inner_code) or "", + error=str(error_msg) or "", + ) + + # Ensure the fixed code is properly stripped and doesn't include markers + fixed_inner_code = result.fixed_code.strip() + if fixed_inner_code.startswith('#') and 'code start' in fixed_inner_code: + # If LLM included markers in response, extract only inner code + inner_match = re.search(r'#\s+\w+\s+code\s+start\s*\n([\s\S]*?)#\s+\w+\s+code\s+end', fixed_inner_code) + if inner_match: + fixed_inner_code = inner_match.group(1).strip() + + # Reconstruct the block with fixed code + fixed_block = f"{start_marker}\n\n{fixed_inner_code}\n\n{end_marker}" + + # Replace the original block with the fixed block in the full code + result_code = result_code.replace(block_code, fixed_block) + logger.log_message(f"Fixed {agent_name} block successfully", level=logging.INFO) + + except Exception as e: + # Log the error but continue with other blocks + logger.log_message(f"Error fixing {agent_name} block: {str(e)}", level=logging.ERROR) + continue + + return result_code def get_dataset_context(df): """ @@ -371,12 +346,13 @@ def get_dataset_context(df): context += f" * {col}: {sample_str}\n" return context except Exception as e: + logger.log_message(f"Error generating dataset context: {str(e)}", level=logging.ERROR) return "Could not generate dataset context information." def edit_code_with_dspy(original_code: str, user_prompt: str, dataset_context: str = ""): - thread_lm = MODEL_OBJECTS['claude-sonnet-4-6'] - with dspy.context(lm=thread_lm): - code_editor = dspy.Predict(code_edit) + gemini = dspy.LM("gemini/gemini-2.5-pro-preview-03-25", api_key = os.environ['GEMINI_API_KEY'], max_tokens=2000) + with dspy.context(lm=gemini): + code_editor = dspy.ChainOfThought(code_edit) result = code_editor( dataset_context=dataset_context, @@ -434,9 +410,8 @@ async def execute_code( # Access app state via request app_state = request.app.state session_state = app_state.get_session_state(session_id) - # logger.log_message(f"Session State: {session_state}", level=logging.INFO) - if session_state["datasets"] is None: + if session_state["current_df"] is None: raise HTTPException( status_code=400, detail="No dataset is currently loaded. Please link a dataset before executing code." @@ -447,138 +422,15 @@ async def execute_code( if not code: raise HTTPException(status_code=400, detail="No code provided") - # Get the user_id and chat_id from session state if available - user_id = session_state.get("user_id") - chat_id = session_state.get("chat_id") - message_id = request_data.message_id - - # If message_id was not provided in the request, try to get it from the session state - if message_id is None: - message_id = session_state.get("current_message_id") - else: - # Update the session state with the provided message_id - session_state["current_message_id"] = message_id - - # Get model configuration - model_config = session_state.get("model_config", {}) - model_provider = model_config.get("provider", "") - model_name = model_config.get("model", "") - model_temperature = model_config.get("temperature", 0.0) - model_max_tokens = model_config.get("max_tokens", 0) - - # Get database session - db = get_session() - - # Check if we have an existing execution record for this message - existing_execution = None - if message_id: - try: - existing_execution = db.query(CodeExecution).filter( - CodeExecution.message_id == message_id - ).first() - - except Exception as query_error: - logger.log_message(f"Error querying for existing execution: {str(query_error)}", level=logging.ERROR) - # Continue without existing execution - else: - logger.log_message("No message_id provided in session state", level=logging.WARNING) - # Execute the code with the dataframe from session state - full_output = "" - json_outputs = [] - matplotlib_outputs = [] - is_successful = True - failed_agents = None - error_messages = None - - try: - full_output, json_outputs, matplotlib_outputs = execute_code_from_markdown(code, session_state["datasets"]) - - # Even with "successful" execution, check for agent failures in the output - failed_blocks = identify_error_blocks(code, full_output) - - if failed_blocks: - # We have some failed agents even though no exception was thrown - is_successful = False # Mark as failed if any agent failed - failed_agents = json.dumps([block[0] for block in failed_blocks]) - error_messages = json.dumps({ - block[0]: block[2] for block in failed_blocks - }) - logger.log_message(f"Partial execution failure. Failed agents: {failed_agents}", level=logging.WARNING) - - except Exception as exec_error: - full_output = str(exec_error) - json_outputs = [] - matplotlib_outputs = [] - is_successful = False - - # Identify which agents failed - failed_blocks = identify_error_blocks(code, full_output) - - # Format the failed agents and error messages - if failed_blocks: - failed_agents = json.dumps([block[0] for block in failed_blocks]) - error_messages = json.dumps({ - block[0]: block[2] for block in failed_blocks - }) - logger.log_message(f"Execution threw exception. Failed agents: {failed_agents}", level=logging.ERROR) - - # Don't re-raise the error - we want to capture the error and send it back to the client - # return error details in the response instead - - # Create or update the execution record regardless of success/failure - try: - if existing_execution: - # Update existing record - existing_execution.latest_code = code - existing_execution.is_successful = is_successful - existing_execution.output = full_output - - if not is_successful: - existing_execution.failed_agents = failed_agents - existing_execution.error_messages = error_messages - - db.commit() - else: - # Create new record - logger.log_message(f"Creating new CodeExecution record for message_id: {message_id}", level=logging.INFO) - new_execution = CodeExecution( - message_id=message_id, - chat_id=chat_id, - user_id=user_id, - initial_code=code, - latest_code=code, - is_successful=is_successful, - output=full_output, - model_provider=model_provider, - model_name=model_name, - model_temperature=model_temperature, - model_max_tokens=model_max_tokens, - failed_agents=failed_agents, - error_messages=error_messages - ) - db.add(new_execution) - db.commit() - logger.log_message(f"Successfully created CodeExecution record with ID: {new_execution.execution_id} for message_id: {message_id}", level=logging.INFO) - except Exception as db_error: - db.rollback() - logger.log_message(f"Error saving code execution: {str(db_error)}", level=logging.ERROR) - finally: - db.close() + output, json_outputs = execute_code_from_markdown(code, session_state["current_df"]) # Format plotly outputs for frontend plotly_outputs = [f"```plotly\n{json_output}\n```\n" for json_output in json_outputs] - # Format matplotlib outputs for frontend - matplotlib_chart_outputs = [f"```matplotlib\n{img_base64}\n```\n" for img_base64 in matplotlib_outputs] - - # Include execution status in the response return { - "output": full_output, - "plotly_outputs": plotly_outputs if json_outputs else None, - "matplotlib_outputs": matplotlib_chart_outputs if matplotlib_outputs else None, - "is_successful": is_successful, - "failed_agents": failed_agents + "output": output, + "plotly_outputs": plotly_outputs if json_outputs else None } except Exception as e: logger.log_message(f"Error executing code: {str(e)}", level=logging.ERROR) @@ -612,7 +464,8 @@ async def edit_code( session_state = app_state.get_session_state(session_id) # Get dataset context - dataset_context = get_dataset_context(session_state["datasets"]) + dataset_context = get_dataset_context(session_state["current_df"]) + try: # Use the configured language model with dataset context edited_code = edit_code_with_dspy( @@ -655,72 +508,38 @@ async def fix_code( Dictionary containing the fixed code and information about fixed blocks """ try: - # Add debugging at the start - logger.log_message(f"🔧 /fix endpoint called with session_id: {session_id}", level=logging.INFO) - logger.log_message(f"🔧 Code length: {len(request_data.code) if request_data.code else 0}", level=logging.INFO) - logger.log_message(f"🔧 Error length: {len(request_data.error) if request_data.error else 0}", level=logging.INFO) - # Check if code and error are provided if not request_data.code or not request_data.error: - logger.log_message(f"Error fixing code: Both code and error message are required {request_data.code} {request_data.error}", level=logging.ERROR) raise HTTPException(status_code=400, detail="Both code and error message are required") # Access app state via request app_state = request.app.state session_state = app_state.get_session_state(session_id) - logger.log_message(f"🔧 Session state keys: {list(session_state.keys()) if session_state else 'None'}", level=logging.INFO) - - # Get the user_id from session state if available (for logging/tracking) - user_id = session_state.get("user_id") - logger.log_message(f"Code fix request from user_id: {user_id}, session_id: {session_id}", level=logging.INFO) - # Get dataset context - logger.log_message(f"🔧 Getting dataset context...", level=logging.INFO) - dataset_context = get_dataset_context(session_state["datasets"]) - logger.log_message(f"🔧 Dataset context length: {len(dataset_context)}", level=logging.INFO) + dataset_context = get_dataset_context(session_state["current_df"]) try: - logger.log_message(f"🔧 Calling fix_code_with_dspy...", level=logging.INFO) # Use the code_fix agent to fix the code, with dataset context - fixed_code = await fix_code_with_dspy( + fixed_code = fix_code_with_dspy( request_data.code, request_data.error, - dataset_context, - session_state["datasets"] # Pass the actual datasets + dataset_context ) - logger.log_message(f"🔧 fix_code_with_dspy returned, formatting...", level=logging.INFO) fixed_code = format_code_block(fixed_code) - - logger.log_message(f"Code fix completed successfully for user_id: {user_id}", level=logging.INFO) - logger.log_message(f"🔧 Fixed code length: {len(fixed_code)}", level=logging.INFO) return { "fixed_code": fixed_code, } except Exception as e: - logger.log_message(f"🔧 Error in fix_code_with_dspy: {str(e)}", level=logging.ERROR) # Fallback if DSPy models are not initialized or there's an error - logger.log_message(f"Error with DSPy models for user_id {user_id}: {str(e)}", level=logging.ERROR) - - # Return the actual error details instead of generic message - error_message = str(e) - - # Sanitize sensitive information but keep useful details - if "API key" in error_message.lower(): - error_message = "API configuration error. Please contact support." - elif "timeout" in error_message.lower(): - error_message = "Request timed out. Please try again." - elif "rate limit" in error_message.lower(): - error_message = "Rate limit exceeded. Please wait a moment and try again." - elif len(error_message) > 200: - # Truncate very long error messages - error_message = error_message[:200] + "..." + logger.log_message(f"Error with DSPy models: {str(e)}", level=logging.ERROR) + # Return a helpful error message that doesn't expose implementation details return { "fixed_code": request_data.code, - "error": error_message # Return actual error instead of generic message + "error": "Could not process fix request. Please try again later." } except Exception as e: logger.log_message(f"Error fixing code: {str(e)}", level=logging.ERROR) @@ -757,73 +576,4 @@ async def clean_code( } except Exception as e: logger.log_message(f"Error cleaning code: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=str(e)) - -@router.post("/get-latest-code") -async def get_latest_code( - request_data: GetLatestCodeRequest, - request: Request, - session_id: str = Depends(get_session_id_dependency) -): - """ - Retrieve the latest code for a specific message_id - - Args: - request_data: Body containing message_id - request: FastAPI Request object - session_id: Session identifier - - Returns: - Dictionary containing the latest code and execution status - """ - try: - message_id = request_data.message_id - - if not message_id: - raise HTTPException(status_code=400, detail="Message ID is required") - - # Get database session - db = get_session() - - try: - # Query the database for the latest code execution record (ordered by created_at desc) - logger.log_message(f"Searching for execution records with message_id: {message_id}", level=logging.INFO) - - execution_record = db.query(CodeExecution).filter( - CodeExecution.message_id == message_id - ).order_by(CodeExecution.created_at.desc()).first() - - # Also log total count of records for this message_id - total_records = db.query(CodeExecution).filter( - CodeExecution.message_id == message_id - ).count() - logger.log_message(f"Found {total_records} execution records for message_id: {message_id}", level=logging.INFO) - - if execution_record: - logger.log_message(f"Latest execution record found - success: {execution_record.is_successful}, latest_code length: {len(execution_record.latest_code or '') if execution_record.latest_code else 0}", level=logging.INFO) - - # Return the latest code and execution status - return { - "found": True, - "message_id": message_id, - "latest_code": execution_record.latest_code, - "initial_code": execution_record.initial_code, - "is_successful": execution_record.is_successful, - "failed_agents": execution_record.failed_agents - } - else: - logger.log_message(f"No execution record found for message_id: {message_id}", level=logging.INFO) - return { - "found": False, - "message_id": message_id - } - - except Exception as db_error: - logger.log_message(f"Database error retrieving latest code: {str(db_error)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Database error: {str(db_error)}") - finally: - db.close() - - except Exception as e: - logger.log_message(f"Error retrieving latest code: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=500, detail=str(e)) \ No newline at end of file diff --git a/src/routes/deep_analysis_routes.py b/src/routes/deep_analysis_routes.py deleted file mode 100644 index c9498e1f105608f343957778cd9a4c9f82a51738..0000000000000000000000000000000000000000 --- a/src/routes/deep_analysis_routes.py +++ /dev/null @@ -1,570 +0,0 @@ -import logging -from fastapi import APIRouter, HTTPException, Query, Body -from typing import List, Optional -from datetime import datetime, UTC -from sqlalchemy import desc -import json - -from src.db.init_db import session_factory -from src.db.schemas.models import DeepAnalysisReport -from src.utils.logger import Logger - -from src.schemas.deep_analysis_schema import DeepAnalysisReportCreate, DeepAnalysisReportResponse, DeepAnalysisReportDetailResponse - -# Initialize logger with console logging disabled -logger = Logger("deep_analysis_routes", see_time=True, console_log=False) - -# Initialize router -router = APIRouter(prefix="/deep_analysis", tags=["deep_analysis"]) - -# Routes -@router.post("/reports", response_model=DeepAnalysisReportResponse) -async def create_report(report: DeepAnalysisReportCreate): - """Store a deep analysis report in the database""" - try: - session = session_factory() - - try: - # Calculate duration if not provided - duration_seconds = None - if report.duration_seconds is not None: - duration_seconds = report.duration_seconds - - # Convert any JSON data to strings for storage - summaries = report.summaries - plotly_figures = report.plotly_figures - synthesis = report.synthesis - - if isinstance(summaries, list): - summaries = json.dumps(summaries) - - if isinstance(plotly_figures, list): - # Handle serialization of plotly figures specially - # We'll store references or simplified versions - plotly_figures = json.dumps(plotly_figures) - - if isinstance(synthesis, list): - synthesis = json.dumps(synthesis) - - # Create a summary if not provided - report_summary = report.report_summary - if not report_summary and report.final_conclusion: - # remove the **Conclusion** and **Key Takeaways** from the final conclusion - report_summary = report.final_conclusion.replace("**Conclusion**", "") - # Create a summary from the conclusion (first 200 chars) - report_summary = report.final_conclusion[:200] + "..." if len(report.final_conclusion) > 200 else report.final_conclusion - - now = datetime.now(UTC) - - new_report = DeepAnalysisReport( - report_uuid=report.report_uuid, - user_id=report.user_id, - goal=report.goal, - status=report.status, - start_time=now, - end_time=now, - duration_seconds=duration_seconds, - deep_questions=report.deep_questions, - deep_plan=report.deep_plan, - summaries=summaries, - analysis_code=report.analysis_code, - plotly_figures=plotly_figures, - synthesis=synthesis, - final_conclusion=report.final_conclusion, - html_report=report.html_report, - report_summary=report_summary, - progress_percentage=report.progress_percentage, - # Credit and error tracking - credits_consumed=report.credits_consumed or 0, - error_message=report.error_message, - model_provider=report.model_provider or "anthropic", - model_name=report.model_name or "claude-sonnet-4-6", - total_tokens_used=report.total_tokens_used or 0, - estimated_cost=report.estimated_cost or 0.0, - steps_completed=report.steps_completed, - created_at=now, - updated_at=now - ) - - session.add(new_report) - session.commit() - session.refresh(new_report) - - # Return response with created report data - return { - "report_id": new_report.report_id, - "report_uuid": new_report.report_uuid, - "user_id": new_report.user_id, - "goal": new_report.goal, - "status": new_report.status, - "start_time": new_report.start_time, - "end_time": new_report.end_time, - "duration_seconds": new_report.duration_seconds, - "report_summary": new_report.report_summary, - "created_at": new_report.created_at, - "updated_at": new_report.updated_at - } - - except Exception as e: - session.rollback() - logger.log_message(f"Error creating deep analysis report: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to create report: {str(e)}") - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error creating deep analysis report: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to create report: {str(e)}") - -@router.get("/reports", response_model=List[DeepAnalysisReportResponse]) -async def get_reports( - user_id: Optional[int] = None, - limit: int = Query(10, ge=1, le=100), - offset: int = Query(0, ge=0), - status: Optional[str] = None -): - """Get deep analysis reports, optionally filtered by user_id or status""" - try: - session = session_factory() - - try: - query = session.query(DeepAnalysisReport) - - if user_id is not None: - query = query.filter(DeepAnalysisReport.user_id == user_id) - - if status is not None: - query = query.filter(DeepAnalysisReport.status == status) - - # Order by most recent first - query = query.order_by(desc(DeepAnalysisReport.created_at)) - - reports = query.limit(limit).offset(offset).all() - - return [{ - "report_id": report.report_id, - "report_uuid": report.report_uuid, - "user_id": report.user_id, - "goal": report.goal, - "status": report.status, - "start_time": report.start_time, - "end_time": report.end_time, - "duration_seconds": report.duration_seconds, - "report_summary": report.report_summary, - "created_at": report.created_at, - "updated_at": report.updated_at - } for report in reports] - - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error retrieving deep analysis reports: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve reports: {str(e)}") - -@router.get("/reports/user_historical", response_model=List[DeepAnalysisReportResponse]) -async def get_user_historical_reports(user_id: int, limit: int = Query(50, ge=1, le=100)): - """Get all historical deep analysis reports for a user""" - try: - session = session_factory() - - try: - reports = session.query(DeepAnalysisReport)\ - .filter(DeepAnalysisReport.user_id == user_id)\ - .order_by(desc(DeepAnalysisReport.created_at))\ - .limit(limit)\ - .all() - - return [{ - "report_id": report.report_id, - "report_uuid": report.report_uuid, - "user_id": report.user_id, - "goal": report.goal, - "status": report.status, - "start_time": report.start_time, - "end_time": report.end_time, - "duration_seconds": report.duration_seconds, - "report_summary": report.report_summary, - "created_at": report.created_at, - "updated_at": report.updated_at - } for report in reports] - - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error retrieving user historical reports: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve historical reports: {str(e)}") - -@router.get("/reports/{report_id}", response_model=DeepAnalysisReportDetailResponse) -async def get_report_by_id(report_id: int, user_id: Optional[int] = None): - """Get a specific deep analysis report by ID""" - try: - session = session_factory() - - try: - query = session.query(DeepAnalysisReport).filter(DeepAnalysisReport.report_id == report_id) - - # If user_id provided, ensure the report belongs to that user - if user_id is not None: - query = query.filter(DeepAnalysisReport.user_id == user_id) - - report = query.first() - - if not report: - raise HTTPException(status_code=404, detail=f"Report with ID {report_id} not found") - - # Parse JSON fields - summaries = report.summaries - plotly_figures = report.plotly_figures - synthesis = report.synthesis - - if isinstance(summaries, str): - try: - summaries = json.loads(summaries) - except: - summaries = [] - - if isinstance(plotly_figures, str): - try: - plotly_figures = json.loads(plotly_figures) - except: - plotly_figures = [] - - if isinstance(synthesis, str): - try: - synthesis = json.loads(synthesis) - except: - synthesis = [] - - return { - "report_id": report.report_id, - "report_uuid": report.report_uuid, - "user_id": report.user_id, - "goal": report.goal, - "status": report.status, - "start_time": report.start_time, - "end_time": report.end_time, - "duration_seconds": report.duration_seconds, - "deep_questions": report.deep_questions, - "deep_plan": report.deep_plan, - "summaries": summaries, - "analysis_code": report.analysis_code, - "plotly_figures": plotly_figures, - "synthesis": synthesis, - "final_conclusion": report.final_conclusion, - "html_report": report.html_report, - "report_summary": report.report_summary, - "progress_percentage": report.progress_percentage, - # Credit and error tracking - "credits_consumed": report.credits_consumed, - "error_message": report.error_message, - "model_provider": report.model_provider, - "model_name": report.model_name, - "total_tokens_used": report.total_tokens_used, - "estimated_cost": report.estimated_cost, - "steps_completed": report.steps_completed, - "created_at": report.created_at, - "updated_at": report.updated_at - } - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error retrieving deep analysis report: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve report: {str(e)}") - -@router.get("/reports/uuid/{report_uuid}", response_model=DeepAnalysisReportDetailResponse) -async def get_report_by_uuid(report_uuid: str, user_id: Optional[int] = None): - """Get a specific deep analysis report by UUID""" - try: - session = session_factory() - - try: - query = session.query(DeepAnalysisReport).filter(DeepAnalysisReport.report_uuid == report_uuid) - - # If user_id provided, ensure the report belongs to that user - if user_id is not None: - query = query.filter(DeepAnalysisReport.user_id == user_id) - - report = query.first() - - if not report: - raise HTTPException(status_code=404, detail=f"Report with UUID {report_uuid} not found") - - # Parse JSON fields - summaries = report.summaries - plotly_figures = report.plotly_figures - synthesis = report.synthesis - - if isinstance(summaries, str): - try: - summaries = json.loads(summaries) - except: - summaries = [] - - if isinstance(plotly_figures, str): - try: - plotly_figures = json.loads(plotly_figures) - except: - plotly_figures = [] - - if isinstance(synthesis, str): - try: - synthesis = json.loads(synthesis) - except: - synthesis = [] - - return { - "report_id": report.report_id, - "report_uuid": report.report_uuid, - "user_id": report.user_id, - "goal": report.goal, - "status": report.status, - "start_time": report.start_time, - "end_time": report.end_time, - "duration_seconds": report.duration_seconds, - "deep_questions": report.deep_questions, - "deep_plan": report.deep_plan, - "summaries": summaries, - "analysis_code": report.analysis_code, - "plotly_figures": plotly_figures, - "synthesis": synthesis, - "final_conclusion": report.final_conclusion, - "html_report": report.html_report, - "report_summary": report.report_summary, - "progress_percentage": report.progress_percentage, - # Credit and error tracking - "credits_consumed": report.credits_consumed, - "error_message": report.error_message, - "model_provider": report.model_provider, - "model_name": report.model_name, - "total_tokens_used": report.total_tokens_used, - "estimated_cost": report.estimated_cost, - "steps_completed": report.steps_completed, - "created_at": report.created_at, - "updated_at": report.updated_at - } - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error retrieving deep analysis report: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve report: {str(e)}") - -@router.delete("/reports/{report_id}") -async def delete_report(report_id: int, user_id: Optional[int] = None): - """Delete a deep analysis report""" - try: - session = session_factory() - - try: - query = session.query(DeepAnalysisReport).filter(DeepAnalysisReport.report_id == report_id) - - # If user_id provided, ensure the report belongs to that user - if user_id is not None: - query = query.filter(DeepAnalysisReport.user_id == user_id) - - report = query.first() - - if not report: - raise HTTPException(status_code=404, detail=f"Report with ID {report_id} not found") - - session.delete(report) - session.commit() - - return {"message": f"Report {report_id} deleted successfully"} - - except HTTPException: - raise - except Exception as e: - session.rollback() - logger.log_message(f"Error deleting deep analysis report: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to delete report: {str(e)}") - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error deleting deep analysis report: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to delete report: {str(e)}") - -@router.put("/reports/{report_id}/status", response_model=DeepAnalysisReportResponse) -async def update_report_status(report_id: int, status: str = Body(..., embed=True), user_id: Optional[int] = None): - """Update the status of a deep analysis report""" - try: - if status not in ["pending", "running", "completed", "failed"]: - raise HTTPException(status_code=400, detail="Invalid status value") - - session = session_factory() - - try: - query = session.query(DeepAnalysisReport).filter(DeepAnalysisReport.report_id == report_id) - - # If user_id provided, ensure the report belongs to that user - if user_id is not None: - query = query.filter(DeepAnalysisReport.user_id == user_id) - - report = query.first() - - if not report: - raise HTTPException(status_code=404, detail=f"Report with ID {report_id} not found") - - # Update status and end_time if completed or failed - report.status = status - if status in ["completed", "failed"]: - report.end_time = datetime.now(UTC) - if report.start_time: - # Calculate duration in seconds - report.duration_seconds = int((report.end_time - report.start_time).total_seconds()) - - report.updated_at = datetime.now(UTC) - session.commit() - session.refresh(report) - - return { - "report_id": report.report_id, - "report_uuid": report.report_uuid, - "user_id": report.user_id, - "goal": report.goal, - "status": report.status, - "start_time": report.start_time, - "end_time": report.end_time, - "duration_seconds": report.duration_seconds, - "report_summary": report.report_summary, - "created_at": report.created_at, - "updated_at": report.updated_at - } - - except HTTPException: - raise - except Exception as e: - session.rollback() - logger.log_message(f"Error updating report status: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to update report status: {str(e)}") - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error updating report status: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to update report status: {str(e)}") - -@router.get("/reports/uuid/{report_uuid}/html", response_model=dict) -async def get_html_report(report_uuid: str, user_id: Optional[int] = None): - """Get only the HTML report for a specific analysis by UUID""" - try: - session = session_factory() - - try: - query = session.query(DeepAnalysisReport).filter(DeepAnalysisReport.report_uuid == report_uuid) - - # If user_id provided, ensure the report belongs to that user - if user_id is not None: - query = query.filter(DeepAnalysisReport.user_id == user_id) - - report = query.first() - - if not report: - raise HTTPException(status_code=404, detail=f"Report with UUID {report_uuid} not found") - - if not report.html_report: - # Attempt to generate a new HTML report if data is available - from app import generate_html_report # Import the function from app.py - import json - - # Extract report data and regenerate HTML - data_for_report = { - "goal": report.goal, - "deep_questions": report.deep_questions or "", - "deep_plan": report.deep_plan or "", - "summaries": json.loads(report.summaries) if report.summaries and isinstance(report.summaries, str) else [], - "code": report.analysis_code or "", - "plotly_figs": json.loads(report.plotly_figures) if report.plotly_figures and isinstance(report.plotly_figures, str) else [], - "synthesis": json.loads(report.synthesis) if report.synthesis and isinstance(report.synthesis, str) else [], - "final_conclusion": report.final_conclusion or "" - } - - try: - html_report = generate_html_report(data_for_report) - - # Store the generated report back in the database - report.html_report = html_report - session.commit() - - except Exception as e: - logger.log_message(f"Error regenerating HTML report: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to generate HTML report: {str(e)}") - - # Create a filename with timestamp - timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") - filename = f"deep_analysis_report_{timestamp}.html" - - return { - "html_report": report.html_report, - "filename": filename - } - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error retrieving HTML report: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve HTML report: {str(e)}") - -@router.post("/download_from_db/{report_uuid}") -async def download_report_from_db(report_uuid: str, user_id: Optional[int] = None): - """Download HTML report directly from the database""" - try: - session = session_factory() - - try: - query = session.query(DeepAnalysisReport).filter(DeepAnalysisReport.report_uuid == report_uuid) - - # If user_id provided, ensure the report belongs to that user - if user_id is not None: - query = query.filter(DeepAnalysisReport.user_id == user_id) - - report = query.first() - - if not report: - raise HTTPException(status_code=404, detail=f"Report with UUID {report_uuid} not found") - - if not report.html_report: - raise HTTPException(status_code=404, detail=f"HTML report not found for {report_uuid}") - - # Create a filename with timestamp - from datetime import datetime - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - filename = f"deep_analysis_report_{timestamp}.html" - - from fastapi.responses import StreamingResponse - - # Return as downloadable file - return StreamingResponse( - iter([report.html_report.encode('utf-8')]), - media_type='text/html', - headers={ - 'Content-Disposition': f'attachment; filename="{filename}"', - 'Content-Type': 'text/html; charset=utf-8' - } - ) - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error downloading report from database: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to download report: {str(e)}") \ No newline at end of file diff --git a/src/routes/feedback_routes.py b/src/routes/feedback_routes.py deleted file mode 100644 index 65ee5b5ddd52fe4ddb20219caae901a435919bc8..0000000000000000000000000000000000000000 --- a/src/routes/feedback_routes.py +++ /dev/null @@ -1,190 +0,0 @@ -from fastapi import APIRouter, Depends, HTTPException, Query -from typing import List, Optional -import logging -from src.db.init_db import session_factory -from src.db.schemas.models import Message, MessageFeedback -from src.schemas.chat_schema import MessageFeedbackCreate, MessageFeedbackResponse -from src.managers.chat_manager import ChatManager -from src.utils.logger import Logger -import os -from dotenv import load_dotenv -from datetime import datetime, UTC - -load_dotenv() - -# Initialize logger -logger = Logger("feedback_routes", see_time=True, console_log=False) - -# Initialize router -router = APIRouter(prefix="/feedback", tags=["feedback"]) - -# Initialize chat manager -chat_manager = ChatManager(db_url=os.getenv("DATABASE_URL")) - -@router.post("/message/{message_id}", response_model=MessageFeedbackResponse) -async def create_message_feedback(message_id: int, feedback: MessageFeedbackCreate): - """Create or update feedback for a message""" - session = session_factory() - try: - # Log the incoming request data - logger.log_message(f"Create feedback request for message {message_id}: {feedback.dict()}", level=logging.INFO) - - # Check if message exists - message = session.query(Message).filter(Message.message_id == message_id).first() - if not message: - logger.log_message(f"Message with ID {message_id} not found", level=logging.WARNING) - raise HTTPException(status_code=404, detail=f"Message with ID {message_id} not found") - - # Check if feedback already exists for this message - existing_feedback = session.query(MessageFeedback).filter( - MessageFeedback.message_id == message_id - ).first() - - now = datetime.now(UTC) - - if existing_feedback: - # Log that we're updating existing feedback - logger.log_message(f"Updating existing feedback (ID: {existing_feedback.feedback_id}) for message {message_id}", level=logging.INFO) - - # Update existing feedback - existing_feedback.rating = feedback.rating - - # Only update these fields if they are provided - if feedback.model_name is not None: - existing_feedback.model_name = feedback.model_name - if feedback.model_provider is not None: - existing_feedback.model_provider = feedback.model_provider - if feedback.temperature is not None: - existing_feedback.temperature = feedback.temperature - if feedback.max_tokens is not None: - existing_feedback.max_tokens = feedback.max_tokens - - existing_feedback.updated_at = now - feedback_record = existing_feedback - else: - # Log that we're creating new feedback - logger.log_message(f"Creating new feedback for message {message_id}", level=logging.INFO) - - # Create new feedback - feedback_record = MessageFeedback( - message_id=message_id, - rating=feedback.rating, - model_name=feedback.model_name, - model_provider=feedback.model_provider, - temperature=feedback.temperature, - max_tokens=feedback.max_tokens, - created_at=now, - updated_at=now - ) - session.add(feedback_record) - - # Commit changes to database - session.commit() - - # Refresh to get updated values - session.refresh(feedback_record) - - # Log success - logger.log_message(f"Successfully saved feedback (ID: {feedback_record.feedback_id}) for message {message_id}", level=logging.INFO) - - # Build response object - response_data = { - "feedback_id": feedback_record.feedback_id, - "message_id": feedback_record.message_id, - "rating": feedback_record.rating, - "feedback_comment": None, # This field is mentioned in schema but not in model - "model_name": feedback_record.model_name, - "model_provider": feedback_record.model_provider, - "temperature": feedback_record.temperature, - "max_tokens": feedback_record.max_tokens, - "created_at": feedback_record.created_at.isoformat(), - "updated_at": feedback_record.updated_at.isoformat() - } - - return response_data - except HTTPException: - raise - except Exception as e: - session.rollback() - logger.log_message(f"Error creating/updating feedback: {str(e)}", level=logging.ERROR) - # Log more detailed error information - import traceback - logger.log_message(f"Traceback: {traceback.format_exc()}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to create/update feedback: {str(e)}") - finally: - session.close() - -@router.get("/message/{message_id}", response_model=MessageFeedbackResponse) -async def get_message_feedback(message_id: int): - """Get feedback for a specific message""" - session = session_factory() - try: - # Check if feedback exists for this message - feedback = session.query(MessageFeedback).filter( - MessageFeedback.message_id == message_id - ).first() - - if not feedback: - raise HTTPException(status_code=404, detail=f"No feedback found for message with ID {message_id}") - - # Safely handle feedback_comment which might be None - feedback_comment = feedback.feedback_comment if hasattr(feedback, 'feedback_comment') else None - - return { - "feedback_id": feedback.feedback_id, - "message_id": feedback.message_id, - "rating": feedback.rating, - "feedback_comment": feedback_comment, - "model_name": feedback.model_name, - "model_provider": feedback.model_provider, - "temperature": feedback.temperature, - "max_tokens": feedback.max_tokens, - "created_at": feedback.created_at.isoformat(), - "updated_at": feedback.updated_at.isoformat() - } - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error retrieving feedback: {str(e)}", level=logging.ERROR) - # Log more detailed error information - import traceback - logger.log_message(f"Traceback: {traceback.format_exc()}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve feedback: {str(e)}") - finally: - session.close() - -@router.get("/chat/{chat_id}", response_model=List[MessageFeedbackResponse]) -async def get_chat_feedback(chat_id: int): - """Get all feedback for messages in a specific chat""" - session = session_factory() - try: - # Query all feedback for messages in this chat - feedback_records = session.query(MessageFeedback).join( - Message, Message.message_id == MessageFeedback.message_id - ).filter( - Message.chat_id == chat_id - ).all() - - if not feedback_records: - return [] - - return [{ - "feedback_id": feedback.feedback_id, - "message_id": feedback.message_id, - "rating": feedback.rating, - "feedback_comment": feedback.feedback_comment if hasattr(feedback, 'feedback_comment') else None, - "model_name": feedback.model_name, - "model_provider": feedback.model_provider, - "temperature": feedback.temperature, - "max_tokens": feedback.max_tokens, - "created_at": feedback.created_at.isoformat(), - "updated_at": feedback.updated_at.isoformat() - } for feedback in feedback_records] - except Exception as e: - logger.log_message(f"Error retrieving chat feedback: {str(e)}", level=logging.ERROR) - # Log detailed error information - import traceback - logger.log_message(f"Traceback: {traceback.format_exc()}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve chat feedback: {str(e)}") - finally: - session.close() \ No newline at end of file diff --git a/src/routes/session_routes.py b/src/routes/session_routes.py index a76122d3e3bfbca49aa7584b3cd1f8d14e591cb5..9526226cc27170aecbc1d4606a44a7f460e297e4 100644 --- a/src/routes/session_routes.py +++ b/src/routes/session_routes.py @@ -1,81 +1,24 @@ import io import logging import json -import re import os from io import StringIO -from typing import Optional, List, Dict -import random +from typing import Optional, List + import pandas as pd from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile from fastapi.security import APIKeyHeader +from pydantic import BaseModel -import numpy as np from src.managers.session_manager import get_session_id -from src.schemas.model_settings_schema import ModelSettings +from src.schemas.model_settings import ModelSettings from src.utils.logger import Logger -from pydantic import BaseModel -from fastapi.responses import JSONResponse -# data context is for excelsheets with multiple sheets and dataset_descrp is for single sheet or csv -from src.agents.agents import data_context_gen, dataset_description_agent -from src.utils.model_registry import MODEL_OBJECTS, mid_lm -from src.utils.dataset_description_generator import generate_dataset_description +from src.agents.agents import dataset_description_agent import dspy -import re -# from fastapi.responses import JSONResponse -import time -# Try to import chardet, but make it optional -try: - import chardet - HAS_CHARDET = True -except ImportError: - HAS_CHARDET = False - logger_temp = Logger("session_routes", see_time=False, console_log=False) - logger_temp.log_message("chardet not installed, encoding detection will be limited", level=logging.WARNING) logger = Logger("session_routes", see_time=False, console_log=False) - -def apply_model_safeguards(model_name: str, provider: str, temperature: float, max_tokens: int) -> dict: - """Apply model-specific safeguards for temperature and max_tokens based on official API limits""" - model_str = str(model_name).lower() - provider_str = str(provider).lower() - - safe_temp = min(1.0, max(0.0, float(temperature))) - safe_max_tokens = max_tokens - - # O-series: temp MUST be 1.0 - if ('o1' in model_str or 'o3' in model_str) and provider_str == 'openai': - safe_temp = 1.0 - safe_max_tokens = min(max_tokens, 100_000) - # GPT-5 series - elif 'gpt-5' in model_str and provider_str == 'openai': - safe_temp = 1.0 - safe_max_tokens = min(max_tokens, 16_000) - # GPT-4 series - elif 'gpt-4' in model_str and provider_str == 'openai': - safe_max_tokens = min(max_tokens, 4_096) - # Anthropic: Sonnet 4/3.7/Opus 4 = 64K, others = 8K - elif provider_str == 'anthropic': - if any(x in model_str for x in ['sonnet-4', 'sonnet-3-7', 'opus-4']): - safe_max_tokens = min(max_tokens, 64_000) - else: - safe_max_tokens = min(max_tokens, 8_192) - # Groq: 32K - elif provider_str == 'groq': - safe_max_tokens = min(max_tokens, 32_768) - # Gemini: 2.5 series = 65K, others = 8K - elif provider_str == 'gemini': - if '2.5' in model_str or '2-5' in model_str: - safe_max_tokens = min(max_tokens, 65_535) - else: - safe_max_tokens = min(max_tokens, 8_192) - else: - safe_max_tokens = min(max_tokens, 4_096) - - return {"temperature": safe_temp, "max_tokens": safe_max_tokens} - # Add session header for dependency X_SESSION_ID = APIKeyHeader(name="X-Session-ID", auto_error=False) @@ -91,12 +34,6 @@ async def get_session_id_dependency(request: Request): app_state = get_app_state(request) return await get_session_id(request, app_state._session_manager) -# Define a model for message tracking -class MessageInfo(BaseModel): - chat_id: Optional[int] = None - message_id: Optional[int] = None - user_id: Optional[int] = None - # Define a model for reset session request class ResetSessionRequest(BaseModel): name: Optional[str] = None @@ -107,12 +44,6 @@ class ResetSessionRequest(BaseModel): class ExcelSheetsResponse(BaseModel): sheets: List[str] -class InitializeSessionRequest(BaseModel): - session_id: str - user_id: int # Changed to int to match your user_id type - user_email: str - user_name: str - @router.post("/api/excel-sheets") async def get_excel_sheets( file: UploadFile = File(...), @@ -130,39 +61,35 @@ async def get_excel_sheets( # Get sheet names sheet_names = excel_file.sheet_names + # Log the sheets found + logger.log_message(f"Found {len(sheet_names)} sheets in Excel file: {', '.join(sheet_names)}", level=logging.INFO) + # Return the sheet names return {"sheets": sheet_names} except Exception as e: logger.log_message(f"Error getting Excel sheets: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=400, detail=f"Error reading Excel file: {str(e)}") - - - @router.post("/upload_excel") async def upload_excel( file: UploadFile = File(...), name: str = Form(...), description: str = Form(...), - selected_sheets: Optional[str] = Form(None), # JSON array of strings - fill_nulls: bool = Form(True), # NEW: Fill null values - convert_types: bool = Form(True), # NEW: Convert data types + sheet_name: str = Form(...), app_state = Depends(get_app_state), session_id: str = Depends(get_session_id_dependency), request: Request = None ): """Upload and process an Excel file with a specific sheet""" try: - logger.log_message(f"Excel upload: fill_nulls={fill_nulls}, convert_types={convert_types}", level=logging.INFO) - # Log the incoming request details - # logger.log_message(f"Excel upload request for session {session_id}: name='{name}', description='{description}', sheet='{sheet_name}'", level=logging.INFO) + logger.log_message(f"Excel upload request for session {session_id}: name='{name}', description='{description}', sheet='{sheet_name}'", level=logging.INFO) # Check if we need to force a complete session reset before upload force_refresh = request.headers.get("X-Force-Refresh") == "true" if request else False if force_refresh: - # logger.log_message(f"Force refresh requested for session {session_id} before Excel upload", level=logging.INFO) + logger.log_message(f"Force refresh requested for session {session_id} before Excel upload", level=logging.INFO) # Reset the session but don't completely wipe it, so we maintain user association app_state.reset_session_to_default(session_id) @@ -170,251 +97,90 @@ async def upload_excel( contents = await file.read() try: - # Load Excel file to get all sheet names - excel_file = pd.ExcelFile(io.BytesIO(contents)) - sheet_names = excel_file.sheet_names - - # Parse selected sheets if provided; else use all sheets - target_sheets = sheet_names - if selected_sheets: - try: - sel = json.loads(selected_sheets) - if isinstance(sel, list): - target_sheets = [s for s in sheet_names if s in sel] - except Exception: - pass - - datasets = {} - processed_sheets = [] + # Read the specific sheet with basic preprocessing + excel_df = pd.read_excel(io.BytesIO(contents), sheet_name=sheet_name) - for sheet_name in target_sheets: - try: - # Read each sheet - sheet_df = pd.read_excel(io.BytesIO(contents), sheet_name=sheet_name) - sheet_df.replace({np.nan: None, np.inf: None, -np.inf: None}, inplace=True) - - # Preprocessing steps - # 1. Drop empty rows and columns - sheet_df.dropna(how='all', inplace=True) - sheet_df.dropna(how='all', axis=1, inplace=True) - - # 2. Clean column names - sheet_df.columns = sheet_df.columns.str.strip() - - # 3. Skip empty sheets - if sheet_df.empty: - continue - - # Register each sheet in DuckDB with a clean table name - clean_sheet_name = clean_dataset_name(sheet_name) - datasets[clean_sheet_name] = sheet_df - - processed_sheets.append(clean_sheet_name) - - except Exception as e: - logger.log_message(f"Error processing sheet '{sheet_name}': {str(e)}", level=logging.WARNING) - continue + # Preprocessing steps + # 1. Drop empty rows and columns + excel_df.dropna(how='all', inplace=True) # Remove empty rows + excel_df.dropna(how='all', axis=1, inplace=True) # Remove empty columns - if not processed_sheets: - raise HTTPException(status_code=400, detail="No valid sheets found in Excel file") + # 2. Clean column names + excel_df.columns = excel_df.columns.str.strip() # Remove extra spaces - # Update the session description (no primary dataset needed) - desc = description - app_state.update_session_dataset(session_id, datasets, processed_sheets, desc) + # 3. Convert Excel data to CSV with UTF-8-sig encoding + csv_buffer = io.StringIO() + excel_df.to_csv(csv_buffer, index=False, encoding='utf-8-sig') + csv_buffer.seek(0) - logger.log_message(f"Processed Excel file with {len(processed_sheets)} sheets: {', '.join(processed_sheets)}", level=logging.INFO) + # Read the processed CSV back into a dataframe + new_df = pd.read_csv(csv_buffer) - return { - "message": "Excel file processed successfully", - "session_id": session_id, - "sheets_processed": processed_sheets, - "total_sheets": len(processed_sheets) - } + # Log some info about the processed data + logger.log_message(f"Processed Excel sheet '{sheet_name}' into dataframe with {len(new_df)} rows and {len(new_df.columns)} columns", level=logging.INFO) except Exception as e: logger.log_message(f"Error processing Excel file: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=400, detail=f"Error processing Excel file: {str(e)}") - + + # Update the dataset description to include sheet name + desc = f"{name} Dataset (from Excel sheet '{sheet_name}'): {description}" + + logger.log_message(f"Updating session dataset with Excel data and description: '{desc}'", level=logging.INFO) + app_state.update_session_dataset(session_id, new_df, name, desc) + + # Log the final state + session_state = app_state.get_session_state(session_id) + logger.log_message(f"Session dataset updated with Excel data and description: '{session_state.get('description')}'", level=logging.INFO) + + return {"message": "Excel file processed successfully", "session_id": session_id, "sheet": sheet_name} except Exception as e: logger.log_message(f"Error in upload_excel: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=400, detail=str(e)) - -def clean_dataset_name(name: str) -> str: - """ - Clean dataset name to be a safe Python identifier. - Removes all characters that would cause issues in Python code execution. - """ - if not name: - return "dataset" - - # Convert to string and strip whitespace - name = str(name).strip() - - # Replace spaces and common separators with underscores - name = re.sub(r'[\s\-\.]+', '_', name) - - # Remove all non-alphanumeric characters except underscores - name = re.sub(r'[^a-zA-Z0-9_]', '', name) - - # Remove multiple consecutive underscores - name = re.sub(r'_+', '_', name) - - # Remove leading/trailing underscores - name = name.strip('_') - - # Ensure it starts with a letter or underscore (Python identifier rule) - if name and not re.match(r'^[a-zA-Z_]', name): - name = f"dataset_{name}" - - # If empty after cleaning, use default - if not name: - name = "dataset" - - # Limit length to 30 characters - if len(name) > 30: - name = name[:30] - - # Ensure it's still a valid identifier after truncation - if not re.match(r'^[a-zA-Z_]', name): - name = f"dataset_{name}" - - - - return name @router.post("/upload_dataframe") async def upload_dataframe( file: UploadFile = File(...), name: str = Form(...), description: str = Form(...), - columns:List[str] = Form(...), - fill_nulls: bool = Form(True), # NEW: Fill null values - convert_types: bool = Form(True), # NEW: Convert data types app_state = Depends(get_app_state), session_id: str = Depends(get_session_id_dependency), request: Request = None ): try: - logger.log_message(f"CSV upload: fill_nulls={fill_nulls}, convert_types={convert_types}", level=logging.INFO) - # Log the incoming request details logger.log_message(f"Upload request for session {session_id}: name='{name}', description='{description}'", level=logging.INFO) # Check if we need to force a complete session reset before upload force_refresh = request.headers.get("X-Force-Refresh") == "true" if request else False - # Log session state BEFORE any changes - session_state_before = app_state.get_session_state(session_id) - datasets_before = session_state_before.get("datasets", {}) - logger.log_message(f"Session state BEFORE upload - datasets: {list(datasets_before.keys())}", level=logging.INFO) - if force_refresh: - logger.log_message(f"Force refresh requested for session {session_id} before CSV upload", level=logging.INFO) + logger.log_message(f"Force refresh requested for session {session_id} before upload", level=logging.INFO) # Reset the session but don't completely wipe it, so we maintain user association app_state.reset_session_to_default(session_id) - - # Log session state AFTER reset - session_state_after_reset = app_state.get_session_state(session_id) - datasets_after_reset = session_state_after_reset.get("datasets", {}) - logger.log_message(f"Session state AFTER reset - datasets: {list(datasets_after_reset.keys())}", level=logging.INFO) - # Clean and validate the name - name = clean_dataset_name(name) - - # Validate name length and create safe variable name - if len(name) > 30: - name = name[:30] - - # Ensure it's a safe Python identifier - - - # Read and process the CSV file - content = await file.read() - new_df = None - last_exception = None - - # Try encodings with delimiter auto-detection (with chardet first) - encodings_to_try = [ - 'utf-8', 'utf-8-sig', 'latin-1', 'cp1252', 'ascii', - 'iso-8859-1', 'iso-8859-2', 'iso-8859-3', 'iso-8859-4', 'iso-8859-5', - 'iso-8859-6', 'iso-8859-7', 'iso-8859-8', 'iso-8859-9', 'iso-8859-15', - 'cp1250', 'cp1251', 'cp1254', 'cp1255', 'cp1256', 'cp1257', - 'cp932', 'shift_jis', 'euc-jp', 'euc-kr', - 'gb2312', 'gbk', 'gb18030', 'big5', 'mac-roman', - 'koi8-r', 'koi8-u' - ] - - - # Try to detect encoding using chardet if available - if HAS_CHARDET: - try: - detected = chardet.detect(content[:100000]) - if detected and detected.get('encoding') and detected.get('confidence', 0) > 0.7: - detected_encoding = detected['encoding'] - if detected_encoding not in encodings_to_try: - encodings_to_try.insert(0, detected_encoding) - logger.log_message(f"Detected encoding: {detected_encoding} (confidence: {detected['confidence']:.2f})", level=logging.INFO) - except Exception: - pass - - delimiters_to_try = [',', ';', '\t', '|', ':', ' '] - - for encoding in encodings_to_try: + # Now process the new file + contents = await file.read() + try: + new_df = pd.read_csv(io.BytesIO(contents), encoding='utf-8') + except Exception as e: try: - csv_content = content.decode(encoding) - sample = csv_content[:1024] - try: - import csv as _csv - dialect = _csv.Sniffer().sniff(sample, delimiters=delimiters_to_try) - delimiter = dialect.delimiter - new_df = pd.read_csv(io.StringIO(csv_content), sep=delimiter, engine='python')[columns] - except Exception: - # Fallback to pandas automatic detection - try: - new_df = pd.read_csv(io.StringIO(csv_content), sep=None, engine='python')[columns] - except Exception: - # Final fallback: brute-force common delimiters - for d in delimiters_to_try: - try: - new_df = pd.read_csv(io.StringIO(csv_content), sep=d, engine='python')[columns] - break - except Exception: - new_df = None - if new_df is not None: - new_df.replace({np.nan: None, np.inf: None, -np.inf: None}, inplace=True) - logger.log_message(f"Successfully read CSV with encoding: {encoding}", level=logging.INFO) - break + new_df = pd.read_csv(io.BytesIO(contents), encoding='unicode_escape') except Exception as e: - last_exception = e - logger.log_message(f"Failed to read CSV with encoding {encoding}: {str(e)}", level=logging.WARNING) - continue - - if new_df is None: - raise HTTPException(status_code=400, detail=f"Error reading file with tried encodings: {encodings_to_try}. Last error: {str(last_exception)}") - - # Format the description - desc = f" exact_python_name: `{name}` Dataset: {description}" - - # Create datasets dictionary with the new dataset - datasets = {name: new_df} - - # Update the session with the new dataset (this will replace any existing datasets) but not update desc, as that is passed already - app_state.update_session_dataset(session_id, datasets, [name], desc, pre_generated=True) - - # Log session state AFTER upload - session_state_after_upload = app_state.get_session_state(session_id) - datasets_after_upload = session_state_after_upload.get("datasets", {}) - logger.log_message(f"Session state AFTER upload - datasets: {list(datasets_after_upload.keys())}", level=logging.INFO) + try: + new_df = pd.read_csv(io.BytesIO(contents), encoding='ISO-8859-1') + except Exception as e: + raise HTTPException(status_code=400, detail=f"Error reading file: {str(e)}") + desc = f"{name} Dataset: {description}" - logger.log_message(f"Successfully uploaded dataset '{name}' for session {session_id}", level=logging.INFO) + logger.log_message(f"Updating session dataset with description: '{desc}'", level=logging.INFO) + app_state.update_session_dataset(session_id, new_df, name, desc) - return JSONResponse(content=sanitize_json({ - "message": "Dataframe uploaded successfully", - "session_id": session_id, - "rows": int(new_df.shape[0]), - "columns": int(new_df.shape[1]) - })) + # Log the final state + session_state = app_state.get_session_state(session_id) + logger.log_message(f"Session dataset updated with description: '{session_state.get('description')}'", level=logging.INFO) + return {"message": "Dataframe uploaded successfully", "session_id": session_id} except Exception as e: logger.log_message(f"Error in upload_dataframe: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=400, detail=str(e)) @@ -439,28 +205,16 @@ async def update_model_settings( # Get session state to update model config session_state = app_state.get_session_state(session_id) - - # Apply model-specific safeguards (temperature + max_tokens) - safe_params = apply_model_safeguards( - model_name=settings.model, - provider=settings.provider, - temperature=settings.temperature, - max_tokens=settings.max_tokens - ) - - # Create the model config with safe parameters + + # Create the model config model_config = { "provider": settings.provider, "model": settings.model, "api_key": settings.api_key, - "temperature": safe_params["temperature"], - "max_tokens": safe_params["max_tokens"] + "temperature": settings.temperature, + "max_tokens": settings.max_tokens } - - # Create the model config - - # Update only the session's model config session_state["model_config"] = model_config @@ -471,15 +225,46 @@ async def update_model_settings( app_state._session_manager._app_model_config = model_config # Create the LM instance to test the configuration, but don't set it globally - lm = MODEL_OBJECTS[str(settings.model)] - - + import dspy + + if settings.provider.lower() == "groq": + logger.log_message(f"Groq Model: {settings.model}", level=logging.INFO) + lm = dspy.GROQ( + model=settings.model, + api_key=settings.api_key, + temperature=settings.temperature, + max_tokens=settings.max_tokens + ) + elif settings.provider.lower() == "anthropic": + logger.log_message(f"Anthropic Model: {settings.model}", level=logging.INFO) + lm = dspy.LM( + model=settings.model, + api_key=settings.api_key, + temperature=settings.temperature, + max_tokens=settings.max_tokens + ) + elif settings.provider.lower() == "gemini": + logger.log_message(f"Gemini Model: {settings.model}, API Key: {settings.api_key}, Temperature: {settings.temperature}, Max Tokens: {settings.max_tokens}", level=logging.INFO) + lm = dspy.LM( + model=f"gemini/{settings.model}", + api_key=settings.api_key, + temperature=settings.temperature, + max_tokens=settings.max_tokens + ) + else: # OpenAI is the default + logger.log_message(f"OpenAI Model: {settings.model}", level=logging.INFO) + lm = dspy.LM( + model=settings.model, + api_key=settings.api_key, + temperature=settings.temperature, + max_tokens=settings.max_tokens + ) # Test the model configuration without setting it globally try: - # resp = lm("Hello, are you working?") - # logger.log_message(f"Model Response: {resp}", level=logging.INFO) + resp = lm("Hello, are you working?") + logger.log_message(f"Model Response: {resp}", level=logging.INFO) # REMOVED: dspy.configure(lm=lm) - no longer set globally return {"message": "Model settings updated successfully"} except Exception as model_error: @@ -518,14 +303,91 @@ async def get_model_settings( # Use values from model_config with fallbacks to defaults return { - "provider": model_config.get("provider", "anthropic"), - "model": model_config.get("model", "claude-sonnet-4-6"), + "provider": model_config.get("provider", "openai"), + "model": model_config.get("model", "gpt-4o-mini"), "hasCustomKey": bool(model_config.get("api_key")) or bool(os.getenv("CUSTOM_API_KEY")), "temperature": model_config.get("temperature", 0.7), "maxTokens": model_config.get("max_tokens", 6000) } - +@router.post("/api/preview-csv") +@router.get("/api/preview-csv") +async def preview_csv(app_state = Depends(get_app_state), session_id: str = Depends(get_session_id_dependency)): + """Preview the dataset stored in the session.""" + try: + # Get the session state to ensure we're using the current dataset + session_state = app_state.get_session_state(session_id) + df = session_state.get("current_df") + + # Handle case where dataset might be missing + if df is None: + logger.log_message(f"Dataset not found in session {session_id}, using default", level=logging.WARNING) + # Create a new default session for this session ID + app_state.reset_session_to_default(session_id) + # Get the session state again + session_state = app_state.get_session_state(session_id) + df = session_state.get("current_df") + + # Replace NaN values with None (which becomes null in JSON) + df = df.where(pd.notna(df), None) + + # Convert columns to appropriate types if necessary + for column in df.columns: + if df[column].dtype == 'object': + # Attempt to convert to boolean if the column contains 'True'/'False' strings + if df[column].isin(['True', 'False']).all(): + df[column] = df[column].astype(bool) + + # Extract name and description if available + name = session_state.get("name", "Dataset") + description = session_state.get("description", "No description available") + + + # Try to get the description from make_data if available + if "make_data" in session_state and session_state["make_data"]: + data_dict = session_state["make_data"] + if "Description" in data_dict: + full_desc = data_dict["Description"] + # Try to parse the description format "{name} Dataset: {description}" + if "Dataset:" in full_desc: + parts = full_desc.split("Dataset:", 1) + extracted_name = parts[0].strip() + extracted_description = parts[1].strip() + + # Only use extracted values if they're meaningful + if extracted_name: + name = extracted_name + if extracted_description and extracted_description != "No description available": + description = extracted_description + + logger.log_message(f"Extracted name: '{name}', description: '{description}'", level=logging.INFO) + else: + # If we can't parse it, use the full description + if full_desc and full_desc != "No description available": + description = full_desc + + # Make sure we're not returning "No description available" if there's a description in the session + if description == "No description available" and session_state.get("description"): + session_desc = session_state.get("description") + # Check if the description is in the format "{name} Dataset: {description}" + if "Dataset:" in session_desc: + parts = session_desc.split("Dataset:", 1) + description = parts[1].strip() + else: + description = session_desc + + # Get rows and convert to dict + preview_data = { + "headers": df.columns.tolist(), + "rows": json.loads(df.head(5).to_json(orient="values")), + "name": name, + "description": description + } + + return preview_data + except Exception as e: + logger.log_message(f"Error in preview_csv: {str(e)}", level=logging.ERROR) + raise HTTPException(status_code=400, detail=str(e)) @router.get("/api/default-dataset") async def get_default_dataset( @@ -539,115 +401,19 @@ async def get_default_dataset( # Get the session state to ensure we're using the default dataset session_state = app_state.get_session_state(session_id) - datasets = session_state["datasets"] - keys = list(datasets.keys()) - if "df" in keys: - df = datasets['df'] - else: - df = pd.read_csv('Housing.csv') - - # Load Housing.csv from the data directory (relative to backend root) - - # Detailed description for the default housing dataset - desc = """ exact_python_name: `Housing_Dataset` Dataset: { - "exact": "Housing_Dataset", - "description": "A dataset containing information about residential properties, including their prices, area, number of bedrooms and bathrooms, number of stories, and various amenities. This dataset is useful for analyzing housing market trends and property valuations.", - "columns": { - "price": { - "type": "integer", - "description": "The selling price of the property in local currency", - "preprocessing": "Direct integer conversion", - "missing_values_handling": "Consider imputing with median price or removing rows with missing values" - }, - "area": { - "type": "integer", - "description": "The total area of the property in square feet", - "preprocessing": "Direct integer conversion", - "missing_values_handling": "Impute with median area or remove rows with missing values" - }, - "bedrooms": { - "type": "integer", - "description": "Number of bedrooms in the property", - "preprocessing": "Direct integer conversion", - "missing_values_handling": "Impute with mode or remove rows with missing values" - }, - "bathrooms": { - "type": "integer", - "description": "Number of bathrooms in the property", - "preprocessing": "Direct integer conversion", - "missing_values_handling": "Impute with mode or remove rows with missing values" - }, - "stories": { - "type": "integer", - "description": "Number of stories in the property", - "preprocessing": "Direct integer conversion", - "missing_values_handling": "Impute with mode or remove rows with missing values" - }, - "mainroad": { - "type": "object", - "description": "Indicates if the property is located on a main road (Yes/No)", - "preprocessing": "Convert to categorical type", - "missing_values_handling": "Impute with mode or create a separate category for missing values" - }, - "guestroom": { - "type": "object", - "description": "Indicates if the property has a guest room (Yes/No)", - "preprocessing": "Convert to categorical type", - "missing_values_handling": "Impute with mode or create a separate category for missing values" - }, - "basement": { - "type": "object", - "description": "Indicates if the property has a basement (Yes/No)", - "preprocessing": "Convert to categorical type", - "missing_values_handling": "Impute with mode or create a separate category for missing values" - }, - "hotwaterheating": { - "type": "object", - "description": "Indicates if the property has hot water heating (Yes/No)", - "preprocessing": "Convert to categorical type", - "missing_values_handling": "Impute with mode or create a separate category for missing values" - }, - "airconditioning": { - "type": "object", - "description": "Indicates if the property has air conditioning (Yes/No)", - "preprocessing": "Convert to categorical type", - "missing_values_handling": "Impute with mode or create a separate category for missing values" - }, - "parking": { - "type": "integer", - "description": "Number of parking spaces available", - "preprocessing": "Direct integer conversion", - "missing_values_handling": "Impute with mode or remove rows with missing values" - }, - "prefarea": { - "type": "object", - "description": "Indicates if the property is located in a preferred area (Yes/No)", - "preprocessing": "Convert to categorical type", - "missing_values_handling": "Impute with mode or create a separate category for missing values" - }, - "furnishingstatus": { - "type": "object", - "description": "Status of furnishing (e.g., furnished, semi-furnished, unfurnished)", - "preprocessing": "Convert to categorical type", - "missing_values_handling": "Impute with mode or create a separate category for missing values" - } - }, - "usage_notes": "When analyzing this dataset, consider the impact of missing values on your analysis. Use appropriate imputation methods to maintain data integrity. Additionally, explore correlations between property features and prices to identify trends in the housing market." -}""" + df = session_state["current_df"] + desc = session_state["description"] + + # Replace NaN values with None (which becomes null in JSON) + df = df.where(pd.notna(df), None) - # Full JSON-safe cleanup (same as CSV preview) - df = df.replace([np.inf, -np.inf], None) # Infs → null - df = df.where(pd.notna(df), None) # NaN → null - df = df.dropna(how="all") # Drop fully-empty rows - df = df.applymap(lambda x: None if isinstance(x, str) and x.strip() == "" else x) - preview_data = { "headers": df.columns.tolist(), - "rows": df.head(10).applymap(to_serializable).values.tolist(), + "rows": df.head(10).values.tolist(), "name": "Housing Dataset", "description": desc } - return JSONResponse(content=sanitize_json(preview_data)) + return preview_data # except Exception as e: # raise HTTPException(status_code=400, detail=str(e)) @@ -656,7 +422,7 @@ async def reset_session( request_data: Optional[ResetSessionRequest] = None, app_state = Depends(get_app_state), session_id: str = Depends(get_session_id_dependency), - names: List[str] = None, + name: str = None, description: str = None ): """Reset session to use default dataset with optional new description""" @@ -684,7 +450,7 @@ async def reset_session( try: session_state = app_state.get_session_state(session_id) session_state["model_config"] = model_config - # logger.log_message(f"Preserved model settings for session {session_id}", level=logging.INFO) + logger.log_message(f"Preserved model settings for session {session_id}", level=logging.INFO) except Exception as e: logger.log_message(f"Failed to restore model settings: {str(e)}", level=logging.ERROR) @@ -694,17 +460,13 @@ async def reset_session( description = request_data.description or description # If name and description are provided, update the dataset description - if names and description: + if name and description: session_state = app_state.get_session_state(session_id) - datasets = session_state["datasets"] + df = session_state["current_df"] desc = f"{description}" - # Ensure datasets is a Dict[str, pd.DataFrame] - if not isinstance(datasets, dict) or not all(isinstance(v, pd.DataFrame) for v in datasets.values()): - - raise HTTPException(status_code=500, detail="Session datasets are not valid DataFrames") # Update the session dataset with the new description - app_state.update_session_dataset(session_id, datasets, names, desc) + app_state.update_session_dataset(session_id, df, name, desc) return { "message": "Session reset to default dataset", @@ -720,80 +482,50 @@ async def reset_session( ) - -@router.post("/generate-description-from-preview") -async def generate_description_from_preview( +@router.post("/create-dataset-description") +async def create_dataset_description( request: dict, app_state = Depends(get_app_state) ): + session_id = request.get("sessionId") + if not session_id: + raise HTTPException(status_code=400, detail="Session ID is required") try: - headers = request.get("headers", []) - rows = request.get("rows", []) - user_description = request.get("description", "") - dataset_name = request.get("name", "Dataset") + # Get the session state to access the dataset + session_state = app_state.get_session_state(session_id) + df = session_state["current_df"] - # Clean the dataset name - dataset_name = clean_dataset_name(dataset_name) + # Get any existing description provided by the user + existing_description = request.get("existingDescription", "") - if not headers or not rows: - raise HTTPException(status_code=400, detail="Headers and rows are required") - - # Convert rows to DataFrame - df = pd.DataFrame(rows, columns=headers) - # Infer data types from the sample data - for col in df.columns: - try: - # Try to convert to numeric - pd.to_numeric(df[col], errors='raise') - df[col] = pd.to_numeric(df[col], errors='coerce') - except: - try: - # Try to convert to datetime (suppress warnings) - import warnings - with warnings.catch_warnings(): - warnings.simplefilter("ignore", UserWarning) - df[col] = pd.to_datetime(df[col], errors='coerce') - # If all values became NaT, it's probably not a date column - if df[col].isna().all(): - df[col] = df[col].astype(str) - except: - # Keep as string - df[col] = df[col].astype(str) + # Convert dataframe to a string representation for the agent + dataset_info = { + "columns": df.columns.tolist(), + "sample": df.head(2).to_dict(), + "stats": df.describe().to_dict() + } - # Build dataset view for description generation - dataset_view = "" - head_data = df.head(3) - columns = [{col: str(head_data[col].dtype)} for col in head_data.columns] - dataset_view += f"exact_table_name={dataset_name}\n:columns:{str(columns)}\n{head_data.to_markdown()}\n" + # Get session-specific model + lm = dspy.LM( + model="gpt-4o-mini", + api_key=os.getenv("OPENAI_API_KEY"), + temperature=0.7, + max_tokens=3000 + ) - # Generate description using AI - with dspy.context(lm=mid_lm): - data_context = dspy.Predict(dataset_description_agent)( - existing_description=user_description, - dataset=dataset_view + # Generate description using session model + with dspy.context(lm=lm): + # If there's an existing description, have the agent improve it + description = dspy.Predict(dataset_description_agent)( + dataset=str(dataset_info), + existing_description=existing_description ) - generated_desc = data_context.description - - # Clean the generated description to ensure it's valid JSON if it's JSON - try: - # Try to parse as JSON to validate it - import json - parsed_desc = json.loads(generated_desc) - # If it's valid JSON, format it properly - cleaned_desc = json.dumps(parsed_desc, indent=2) - except json.JSONDecodeError: - # If it's not JSON, use as-is but clean any problematic characters - cleaned_desc = generated_desc.replace('\\r\\n', '\n').replace('\\n', '\n').replace("\\'", "'") - - # Format the description with exact_python_name - formatted_desc = f" exact_python_name: `{dataset_name}` Dataset: {cleaned_desc}" - - return {"description": formatted_desc} + + return {"description": description.description} except Exception as e: - logger.log_message(f"Failed to generate description from preview: {str(e)}", level=logging.ERROR) raise HTTPException(status_code=500, detail=f"Failed to generate description: {str(e)}") @router.get("/api/session-info") @@ -822,12 +554,14 @@ async def get_session_info( is_custom = True # Also check by checking if we have a dataframe that's different from default - if "datasets" in session_state and session_state["datasets"] is not None: + if "current_df" in session_state and session_state["current_df"] is not None: try: # This is just a basic check - we could make it more sophisticated if needed - key_count = len(session_state["datasets"].keys) - if key_count > 1: - is_custom = True + custom_col_count = len(session_state["current_df"].columns) + if hasattr(session_manager, "_default_df") and session_manager._default_df is not None: + default_col_count = len(session_manager._default_df.columns) + if custom_col_count != default_col_count: + is_custom = True except Exception as e: logger.log_message(f"Error comparing datasets: {str(e)}", level=logging.ERROR) @@ -851,259 +585,4 @@ async def get_session_info( "is_custom_dataset": False, "has_session": False, "error": str(e) - } - -# Add a new route to set the current message ID in the session -@router.post("/set-message-info") -async def set_message_info( - message_info: MessageInfo, - app_state = Depends(get_app_state), - session_id: str = Depends(get_session_id_dependency) -): - """Set the current message ID, chat ID, and user ID in the session""" - try: - # Get the session state - session_state = app_state.get_session_state(session_id) - - # Make a copy of previous values for logging - previous_message_id = session_state.get("current_message_id") - previous_chat_id = session_state.get("chat_id") - previous_user_id = session_state.get("user_id") - - # Update the session with message information - if message_info.message_id is not None: - session_state["current_message_id"] = message_info.message_id - if message_info.chat_id is not None: - session_state["chat_id"] = message_info.chat_id - if message_info.user_id is not None: - session_state["user_id"] = message_info.user_id - - # Get updated values for logging - current_message_id = session_state.get("current_message_id") - current_chat_id = session_state.get("chat_id") - current_user_id = session_state.get("user_id") - - # # Log changes - # logger.log_message( - # f"Message info updated for session {session_id}:\n" - # f" message_id: {previous_message_id} -> {current_message_id}\n" - # f" chat_id: {previous_chat_id} -> {current_chat_id}\n" - # f" user_id: {previous_user_id} -> {current_user_id}", - # level=logging.INFO - # ) - - # Verify session state was updated - updated_session_state = app_state.get_session_state(session_id) - # logger.log_message( - # f"Verified session state after update:\n" - # f" current_message_id: {updated_session_state.get('current_message_id')}\n" - # f" chat_id: {updated_session_state.get('chat_id')}\n" - # f" user_id: {updated_session_state.get('user_id')}", - # level=logging.INFO - # ) - - return { - "success": True, - "session_id": session_id, - "message_id": current_message_id, - "chat_id": current_chat_id, - "user_id": current_user_id - } - except Exception as e: - logger.log_message(f"Error setting message info: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=str(e)) - -def to_serializable(val): - if pd.isna(val): - return None - if isinstance(val, (np.generic,)): - return val.item() - if hasattr(val, "isoformat"): # Handle datetimes - return val.isoformat() - return val - -def sanitize_json(obj): - import math - if isinstance(obj, float): - return None if (math.isnan(obj) or math.isinf(obj)) else obj - if isinstance(obj, dict): - return {k: sanitize_json(v) for k, v in obj.items()} - if isinstance(obj, list): - return [sanitize_json(i) for i in obj] - return obj - - -@router.post("/preview-csv-upload") -async def preview_csv_upload( - file: UploadFile = File(...), -): - """Preview CSV file without modifying session""" - try: - content = await file.read() - - # Try encodings with delimiter auto-detection (with chardet first) - encodings_to_try = [ - 'utf-8', 'utf-8-sig', 'latin-1', 'cp1252', 'ascii', - 'iso-8859-1', 'iso-8859-2', 'iso-8859-3', 'iso-8859-4', 'iso-8859-5', - 'iso-8859-6', 'iso-8859-7', 'iso-8859-8', 'iso-8859-9', 'iso-8859-15', - 'cp1250', 'cp1251', 'cp1254', 'cp1255', 'cp1256', 'cp1257', - 'cp932', 'shift_jis', 'euc-jp', 'euc-kr', - 'gb2312', 'gbk', 'gb18030', 'big5', 'mac-roman', - 'koi8-r', 'koi8-u' - ] - - - # Try to detect encoding using chardet if available - if HAS_CHARDET: - try: - detected = chardet.detect(content[:100000]) - if detected and detected.get('encoding') and detected.get('confidence', 0) > 0.7: - detected_encoding = detected['encoding'] - if detected_encoding not in encodings_to_try: - encodings_to_try.insert(0, detected_encoding) - logger.log_message(f"Preview detected encoding: {detected_encoding} (confidence: {detected['confidence']:.2f})", level=logging.INFO) - except Exception: - pass - - delimiters_to_try = [',', ';', '\t', '|', ':', ' '] - new_df = None - last_exception = None - - for encoding in encodings_to_try: - try: - csv_content = content.decode(encoding) - sample = csv_content[:4096] - try: - import csv as _csv - dialect = _csv.Sniffer().sniff(sample, delimiters=delimiters_to_try) - delimiter = dialect.delimiter - new_df = pd.read_csv(io.StringIO(csv_content), sep=delimiter, engine='python') - except Exception: - # Fallback to pandas automatic detection - try: - new_df = pd.read_csv(io.StringIO(csv_content), sep=None, engine='python') - except Exception: - # Final fallback: brute-force common delimiters - for d in delimiters_to_try: - try: - new_df = pd.read_csv(io.StringIO(csv_content), sep=d, engine='python') - break - except Exception: - new_df = None - if new_df is not None: - logger.log_message(f"Successfully read CSV preview with encoding: {encoding}", level=logging.INFO) - break - except Exception as e: - last_exception = e - logger.log_message(f"Failed to read CSV preview with encoding {encoding}: {str(e)}", level=logging.WARNING) - continue - - if new_df is None: - raise HTTPException(status_code=400, detail=f"Error reading file with tried encodings: {encodings_to_try}. Last error: {str(last_exception)}") - - # Clean and validate the name - name = file.filename.replace('.csv', '').replace(' ', '_').lower().strip() - - # Validate name length and create safe variable name - name = clean_dataset_name(name) - - # Ensure it's a safe Python identifier - - - # Format the description - desc = f" exact_python_name: `{name}` Dataset: {file.filename}" - - # Create datasets dictionary with the new dataset - - - # Update the session with the new dataset (this will replace any existing datasets) - - logger.log_message(f"Successfully previewed dataset '{name}'", level=logging.INFO) - - # Inline this in your CSV preview endpoint right before returning JSONResponse - - # df is your DataFrame built from the uploaded CSV - - # JSON-safe cleanup (no separate helper) - new_df = new_df.replace([np.inf, -np.inf], None) # Infs → null - new_df = new_df.where(pd.notna(new_df), None) # NaN → null - new_df = new_df.dropna(how="all") # Drop fully-empty rows - new_df = new_df.applymap(lambda x: None if isinstance(x, str) and x.strip() == "" else x) - - preview_rows = new_df.head(10).applymap(to_serializable).values.tolist() - - - # Limit preview rows - payload = { - "headers": new_df.columns.tolist(), - "rows": preview_rows, - "name": name, - "description": desc - } - return JSONResponse(content=sanitize_json(payload)) - - except Exception as e: - logger.log_message(f"Error in preview_csv_upload: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=400, detail=str(e)) - -@router.post("/generate-session") -async def generate_session(): - """Generate a new session ID and initialize it with default dataset""" - try: - import uuid - session_id = str(uuid.uuid4()) - - # Initialize the session with default dataset - # This will be handled by the first request to any endpoint that uses get_session_id_dependency - - logger.log_message(f"Generated new session ID: {session_id}", level=logging.INFO) - - return { - "session_id": session_id, - "message": "Session created successfully" - } - except Exception as e: - logger.log_message(f"Error generating session: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to generate session: {str(e)}") - -@router.post("/initialize-session") -async def initialize_session( - request: InitializeSessionRequest, - app_state = Depends(get_app_state) -): - """Initialize session immediately after auth or auto-generation""" - try: - session_id = request.session_id - user_id = request.user_id - user_email = request.user_email - user_name = request.user_name - - logger.log_message(f"🔐 Initializing session for user {user_id}: {session_id}", level=logging.INFO) - - # Create or get session state - session_state = app_state.get_session_state(session_id) - - # Associate user with session (only if user_id > 0) - if user_id > 0: - app_state.set_session_user( - session_id=session_id, - user_id=user_id - ) - - # Store additional user info - session_state["user_email"] = user_email - session_state["user_name"] = user_name - session_state["initialized_at"] = time.time() - - logger.log_message(f"✅ Session initialized successfully for user {user_id}", level=logging.INFO) - - return { - "status": "success", - "session_id": session_id, - "message": "Session initialized successfully", - "user_id": user_id - } - - except Exception as e: - logger.log_message(f"❌ Error initializing session: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to initialize session: {str(e)}") + } \ No newline at end of file diff --git a/src/routes/templates_routes.py b/src/routes/templates_routes.py deleted file mode 100644 index 9ccfc4a156f47f5de4e291ebe56c44d696a82857..0000000000000000000000000000000000000000 --- a/src/routes/templates_routes.py +++ /dev/null @@ -1,713 +0,0 @@ -import logging -import os -from fastapi import APIRouter, Depends, HTTPException, Query, Body -from pydantic import BaseModel, Field -from typing import List, Optional, Dict, Any -from datetime import datetime, UTC -from sqlalchemy import desc, func -from sqlalchemy.exc import IntegrityError - -from src.db.init_db import session_factory -from src.db.schemas.models import AgentTemplate, User, UserTemplatePreference -from src.utils.logger import Logger -from src.agents.agents import toggle_user_template_preference -from src.schemas.template_schema import TemplateResponse, UserTemplatePreferenceResponse, TogglePreferenceRequest - -# Initialize logger with console logging disabled -logger = Logger("templates_routes", see_time=True, console_log=False) - -# Initialize router -router = APIRouter(prefix="/templates", tags=["templates"]) - - -def get_global_usage_counts(session, template_ids: List[int] = None) -> Dict[int, int]: - """ - Calculate global usage counts for templates by summing usage_count across all users. - - Args: - session: Database session - template_ids: Optional list of template IDs to filter by. If None, gets all templates. - - Returns: - Dict mapping template_id to global usage count - """ - try: - query = session.query( - UserTemplatePreference.template_id, - func.sum(UserTemplatePreference.usage_count).label('total_usage') - ).group_by(UserTemplatePreference.template_id) - - if template_ids: - query = query.filter(UserTemplatePreference.template_id.in_(template_ids)) - - results = query.all() - - # Convert to dictionary, defaulting to 0 for templates with no usage - usage_dict = {template_id: int(total_usage or 0) for template_id, total_usage in results} - - # If specific template_ids were requested, ensure all are represented - if template_ids: - for template_id in template_ids: - if template_id not in usage_dict: - usage_dict[template_id] = 0 - - return usage_dict - - except Exception as e: - logger.log_message(f"Error calculating global usage counts: {str(e)}", level=logging.ERROR) - return {} - -# Routes -@router.get("/", response_model=List[TemplateResponse]) -async def get_all_templates(variant_type: str = Query(default="all", description="Filter by variant type: 'individual', 'planner', or 'all'")): - """Get all available agent templates with global usage statistics""" - try: - session = session_factory() - - try: - # Single query to get all active templates - templates = session.query(AgentTemplate).filter(AgentTemplate.is_active == True).all() - - # Filter in Python instead of database - if variant_type and variant_type != "all": - if variant_type == "individual": - templates = [t for t in templates if t.variant_type in ['individual', 'both']] - elif variant_type == "planner": - templates = [t for t in templates if t.variant_type in ['planner', 'both']] - - # Get template IDs for usage calculation (only for filtered templates) - template_ids = [template.template_id for template in templates] - - # Calculate global usage counts (only for the templates we're returning) - global_usage = get_global_usage_counts(session, template_ids) - - return [TemplateResponse( - template_id=template.template_id, - template_name=template.template_name, - display_name=template.display_name, - description=template.description, - prompt_template=template.prompt_template, - template_category=template.category, - icon_url=template.icon_url, - is_premium_only=template.is_premium_only, - is_active=template.is_active, - usage_count=global_usage.get(template.template_id, 0), - created_at=template.created_at, - updated_at=template.updated_at - ) for template in templates] - - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error retrieving templates: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve templates: {str(e)}") - -@router.get("/user/{user_id}", response_model=List[UserTemplatePreferenceResponse]) -async def get_user_template_preferences(user_id: int, variant_type: str = Query(default="planner", description="Filter by variant type: 'individual', 'planner', or 'all'")): - """Get all templates with user preferences (enabled/disabled status and usage)""" - try: - session = session_factory() - - try: - # Validate user exists - user = session.query(User).filter(User.user_id == user_id).first() - if not user: - raise HTTPException(status_code=404, detail="User not found") - - # Get templates filtered by variant type (default to planner for modal) - query = session.query(AgentTemplate).filter(AgentTemplate.is_active == True) - - # Filter by variant type - if variant_type and variant_type != "all": - if variant_type == "individual": - query = query.filter(AgentTemplate.variant_type.in_(['individual', 'both'])) - elif variant_type == "planner": - query = query.filter(AgentTemplate.variant_type.in_(['planner', 'both'])) - else: - # Invalid variant_type, default to planner for modal - query = query.filter(AgentTemplate.variant_type.in_(['planner', 'both'])) - - templates = query.all() - - # Get list of default agent names that should be enabled by default - # Use planner variants when filtering for planner, individual variants otherwise - if variant_type == "planner": - default_agent_names = [ - "planner_preprocessing_agent", - "planner_statistical_analytics_agent", - "planner_sk_learn_agent", - "planner_data_viz_agent" - ] - else: - default_agent_names = [ - "preprocessing_agent", - "statistical_analytics_agent", - "sk_learn_agent", - "data_viz_agent" - ] - - result = [] - for template in templates: - # Get user preference for this template if it exists - preference = session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - # Determine if template should be enabled by default - is_default_agent = template.template_name in default_agent_names - default_enabled = is_default_agent # Default agents enabled by default, others disabled - - # Template is enabled by default for default agents, disabled for others - is_enabled = preference.is_enabled if preference else default_enabled - - result.append(UserTemplatePreferenceResponse( - template_id=template.template_id, - template_name=template.template_name, - display_name=template.display_name, - description=template.description, - template_category=template.category, - icon_url=template.icon_url, - is_premium_only=template.is_premium_only, - is_active=template.is_active, - is_enabled=is_enabled, - usage_count=preference.usage_count if preference else 0, - last_used_at=preference.last_used_at if preference else None, - created_at=preference.created_at if preference else None, - updated_at=preference.updated_at if preference else None - )) - - return result - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error retrieving user template preferences: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve user template preferences: {str(e)}") - -@router.get("/user/{user_id}/enabled", response_model=List[UserTemplatePreferenceResponse]) -async def get_user_enabled_templates(user_id: int, variant_type: str = Query(default="planner", description="Filter by variant type: 'individual', 'planner', or 'all'")): - """Get only templates that are enabled for the user (all templates enabled by default)""" - try: - session = session_factory() - - try: - # Validate user exists - user = session.query(User).filter(User.user_id == user_id).first() - if not user: - raise HTTPException(status_code=404, detail="User not found") - - # Get templates filtered by variant type (default to planner for modal) - query = session.query(AgentTemplate).filter(AgentTemplate.is_active == True) - - # Filter by variant type - if variant_type and variant_type != "all": - if variant_type == "individual": - query = query.filter(AgentTemplate.variant_type.in_(['individual', 'both'])) - elif variant_type == "planner": - query = query.filter(AgentTemplate.variant_type.in_(['planner', 'both'])) - else: - # Invalid variant_type, default to planner for modal - query = query.filter(AgentTemplate.variant_type.in_(['planner', 'both'])) - - all_templates = query.all() - - # Get list of default agent names that should be enabled by default - # Use planner variants when filtering for planner, individual variants otherwise - if variant_type == "planner": - default_agent_names = [ - "planner_preprocessing_agent", - "planner_statistical_analytics_agent", - "planner_sk_learn_agent", - "planner_data_viz_agent" - ] - else: - default_agent_names = [ - "preprocessing_agent", - "statistical_analytics_agent", - "sk_learn_agent", - "data_viz_agent" - ] - - result = [] - for template in all_templates: - # Check if user has a preference record for this template - preference = session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - # Determine if template should be enabled by default - is_default_agent = template.template_name in default_agent_names - default_enabled = is_default_agent # Default agents enabled by default, others disabled - - # Template is enabled by default for default agents, disabled for others - is_enabled = preference.is_enabled if preference else default_enabled - - if is_enabled: - result.append(UserTemplatePreferenceResponse( - template_id=template.template_id, - template_name=template.template_name, - display_name=template.display_name, - description=template.description, - template_category=template.category, - icon_url=template.icon_url, - is_premium_only=template.is_premium_only, - is_active=template.is_active, - is_enabled=True, - usage_count=preference.usage_count if preference else 0, - last_used_at=preference.last_used_at if preference else None, - created_at=preference.created_at if preference else None, - updated_at=preference.updated_at if preference else None - )) - - return result - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error retrieving user enabled templates: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve user enabled templates: {str(e)}") - -@router.get("/user/{user_id}/enabled/planner", response_model=List[UserTemplatePreferenceResponse]) -async def get_user_enabled_templates_for_planner(user_id: int): - """Get enabled templates for planner use (max 10 templates)""" - try: - session = session_factory() - - try: - # Validate user exists - user = session.query(User).filter(User.user_id == user_id).first() - if not user: - raise HTTPException(status_code=404, detail="User not found") - - # Get list of default planner agent names that should be enabled by default - default_planner_agent_names = [ - "planner_preprocessing_agent", - "planner_statistical_analytics_agent", - "planner_sk_learn_agent", - "planner_data_viz_agent" - ] - - # Get all active planner variant templates - all_templates = session.query(AgentTemplate).filter( - AgentTemplate.is_active == True, - AgentTemplate.variant_type.in_(['planner', 'both']) - ).all() - - enabled_templates = [] - for template in all_templates: - # Check if user has a preference record for this template - preference = session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - # Determine if template should be enabled by default - is_default_planner_agent = template.template_name in default_planner_agent_names - default_enabled = is_default_planner_agent # Default planner agents enabled by default, others disabled - - # Template is enabled by default for default agents, disabled for others - is_enabled = preference.is_enabled if preference else default_enabled - - if is_enabled: - enabled_templates.append({ - 'template': template, - 'preference': preference, - 'usage_count': preference.usage_count if preference else 0, - 'last_used_at': preference.last_used_at if preference else None - }) - - # Sort by usage (most used first) and limit to 10 - enabled_templates.sort(key=lambda x: (x['usage_count'], x['last_used_at'] or datetime.min.replace(tzinfo=UTC)), reverse=True) - enabled_templates = enabled_templates[:10] - - result = [] - for item in enabled_templates: - template = item['template'] - preference = item['preference'] - - result.append(UserTemplatePreferenceResponse( - template_id=template.template_id, - template_name=template.template_name, - display_name=template.display_name, - description=template.description, - template_category=template.category, - icon_url=template.icon_url, - is_premium_only=template.is_premium_only, - is_active=template.is_active, - is_enabled=True, - usage_count=preference.usage_count if preference else 0, - last_used_at=preference.last_used_at if preference else None, - created_at=preference.created_at if preference else None, - updated_at=preference.updated_at if preference else None - )) - - logger.log_message(f"Retrieved {len(result)} enabled templates for planner for user {user_id}", level=logging.INFO) - return result - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error retrieving planner templates for user {user_id}: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve planner templates: {str(e)}") - -@router.post("/user/{user_id}/template/{template_id}/toggle") -async def toggle_template_preference(user_id: int, template_id: int, request: TogglePreferenceRequest): - """Toggle a user's template preference (enable/disable for planner use)""" - try: - session = session_factory() - - try: - # Validate user exists - user = session.query(User).filter(User.user_id == user_id).first() - if not user: - raise HTTPException(status_code=404, detail="User not found") - - # If trying to disable, check if this would leave user with no enabled templates - if not request.is_enabled: - # Get list of default planner agent names that should be enabled by default - # This function is primarily used by the templates modal which works with planner variants - default_agent_names = [ - "planner_preprocessing_agent", - "planner_statistical_analytics_agent", - "planner_sk_learn_agent", - "planner_data_viz_agent" - ] - - # Get all active planner templates (since this is used by the templates modal) - all_templates = session.query(AgentTemplate).filter( - AgentTemplate.is_active == True, - AgentTemplate.variant_type.in_(['planner', 'both']) - ).all() - - enabled_count = 0 - for template in all_templates: - # Check if user has a preference record for this template - preference = session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - # Determine if template should be enabled by default - is_default_agent = template.template_name in default_agent_names - default_enabled = is_default_agent # Default agents enabled by default, others disabled - - # Template is enabled by default for default agents, disabled for others - is_enabled = preference.is_enabled if preference else default_enabled - - if is_enabled: - enabled_count += 1 - - # If disabling this template would leave the user with 0 enabled templates, reject - if enabled_count <= 1: - raise HTTPException( - status_code=400, - detail="Cannot disable the last active agent. At least one agent must remain active." - ) - - success, message = toggle_user_template_preference( - user_id, template_id, request.is_enabled, session - ) - - if not success: - raise HTTPException(status_code=400, detail=message) - - logger.log_message(f"Toggled template {template_id} for user {user_id}: {message}", level=logging.INFO) - - return {"message": message} - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error toggling template preference: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to toggle template preference: {str(e)}") - -@router.post("/user/{user_id}/bulk-toggle") -async def bulk_toggle_template_preferences(user_id: int, request: dict): - """Bulk toggle multiple template preferences""" - try: - session = session_factory() - - try: - # Validate user exists - user = session.query(User).filter(User.user_id == user_id).first() - if not user: - raise HTTPException(status_code=404, detail="User not found") - - template_preferences = request.get("preferences", []) - if not template_preferences: - raise HTTPException(status_code=400, detail="No preferences provided") - - # Get list of default planner agent names that should be enabled by default - default_planner_agent_names = [ - "planner_preprocessing_agent", - "planner_statistical_analytics_agent", - "planner_sk_learn_agent", - "planner_data_viz_agent" - ] - - # Calculate current enabled count properly (including defaults) - # Focus on planner variants since this is used by the templates modal - all_templates = session.query(AgentTemplate).filter( - AgentTemplate.is_active == True, - AgentTemplate.variant_type.in_(['planner', 'both']) - ).all() - - current_enabled_count = 0 - for template in all_templates: - # Check if user has a preference record for this template - preference = session.query(UserTemplatePreference).filter( - UserTemplatePreference.user_id == user_id, - UserTemplatePreference.template_id == template.template_id - ).first() - - # Determine if template should be enabled by default - is_default_planner_agent = template.template_name in default_planner_agent_names - default_enabled = is_default_planner_agent # Default planner agents enabled by default, others disabled - - # Template is enabled by default for default agents, disabled for others - is_enabled = preference.is_enabled if preference else default_enabled - - if is_enabled: - current_enabled_count += 1 - - # Count how many templates we're trying to enable/disable - enabling_count = sum(1 for pref in template_preferences if pref.get("is_enabled", False)) - disabling_count = sum(1 for pref in template_preferences if not pref.get("is_enabled", False)) - - # Calculate what the new count would be - projected_enabled_count = current_enabled_count + enabling_count - disabling_count - - # Check if the bulk operation would leave user with 0 enabled templates - if projected_enabled_count < 1: - raise HTTPException( - status_code=400, - detail="Cannot disable all agents. At least one agent must remain active." - ) - - results = [] - for pref in template_preferences: - template_id = pref.get("template_id") - is_enabled = pref.get("is_enabled", True) - - if template_id is None: - results.append({"template_id": None, "success": False, "message": "Template ID required"}) - continue - - # Check 10-template limit for enabling - if is_enabled and projected_enabled_count > 10: - results.append({ - "template_id": template_id, - "success": False, - "message": "Cannot enable more than 10 templates for planner use", - "is_enabled": False - }) - continue - - success, message = toggle_user_template_preference( - user_id, template_id, is_enabled, session - ) - - results.append({ - "template_id": template_id, - "success": success, - "message": message, - "is_enabled": is_enabled - }) - - logger.log_message(f"Bulk toggled {len(template_preferences)} templates for user {user_id}", level=logging.INFO) - - return {"results": results} - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error bulk toggling template preferences: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to bulk toggle template preferences: {str(e)}") - -@router.get("/template/{template_id}", response_model=TemplateResponse) -async def get_template(template_id: int): - """Get a specific template by ID with global usage statistics""" - try: - session = session_factory() - - try: - template = session.query(AgentTemplate).filter( - AgentTemplate.template_id == template_id - ).first() - - if not template: - raise HTTPException(status_code=404, detail=f"Template with ID {template_id} not found") - - # Calculate global usage count for this template - global_usage = get_global_usage_counts(session, [template_id]) - - return TemplateResponse( - template_id=template.template_id, - template_name=template.template_name, - display_name=template.display_name, - description=template.description, - prompt_template=template.prompt_template, - template_category=template.category, - icon_url=template.icon_url, - is_premium_only=template.is_premium_only, - is_active=template.is_active, - usage_count=global_usage.get(template_id, 0), # Global usage count - created_at=template.created_at, - updated_at=template.updated_at - ) - - finally: - session.close() - - except HTTPException: - raise - except Exception as e: - logger.log_message(f"Error retrieving template: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve template: {str(e)}") - -@router.get("/categories/list") -async def get_template_categories(): - """Get list of all template categories""" - try: - session = session_factory() - - try: - categories = session.query(AgentTemplate.category).filter( - AgentTemplate.is_active == True, - AgentTemplate.category.isnot(None) - ).distinct().all() - - category_list = [category[0] for category in categories if category[0]] - - return {"categories": category_list} - - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error retrieving template categories: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve template categories: {str(e)}") - -@router.get("/categories") -async def get_templates_by_categories(variant_type: str = Query(default="individual", description="Filter by variant type: 'individual', 'planner', or 'all'")): - """Get all templates grouped by category for frontend template browser with global usage statistics""" - try: - session = session_factory() - - try: - # Get templates filtered by variant type - query = session.query(AgentTemplate).filter(AgentTemplate.is_active == True) - - # Filter by variant type if specified - if variant_type and variant_type != "all": - if variant_type == "individual": - query = query.filter(AgentTemplate.variant_type.in_(['individual', 'both'])) - elif variant_type == "planner": - query = query.filter(AgentTemplate.variant_type.in_(['planner', 'both'])) - else: - # Invalid variant_type, default to individual - query = query.filter(AgentTemplate.variant_type.in_(['individual', 'both'])) - - templates = query.order_by(AgentTemplate.category, AgentTemplate.template_name).all() - - # Get template IDs for usage calculation - template_ids = [template.template_id for template in templates] - - # Calculate global usage counts - global_usage = get_global_usage_counts(session, template_ids) - - # Group templates by category - categories_dict = {} - for template in templates: - category = template.category or "Uncategorized" - if category not in categories_dict: - categories_dict[category] = [] - - categories_dict[category].append({ - "agent_id": template.template_id, # Use template_id as agent_id for compatibility - "agent_name": template.template_name, - "display_name": template.display_name or template.template_name, - "description": template.description, - "prompt_template": template.prompt_template, - "template_category": template.category, - "icon_url": template.icon_url, - "is_premium_only": template.is_premium_only, - "is_active": template.is_active, - "usage_count": global_usage.get(template.template_id, 0), # Global usage count - "created_at": template.created_at.isoformat() if template.created_at else None - }) - - # Convert to list format expected by frontend - result = [] - for category, templates in categories_dict.items(): - result.append({ - "category": category, - "templates": templates - }) - - return result - - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error retrieving templates by categories: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve templates by categories: {str(e)}") - -@router.get("/category/{category}") -async def get_templates_by_category(category: str): - """Get all templates in a specific category with global usage statistics""" - try: - session = session_factory() - - try: - templates = session.query(AgentTemplate).filter( - AgentTemplate.is_active == True, - AgentTemplate.category == category - ).all() - - # Get template IDs for usage calculation - template_ids = [template.template_id for template in templates] - - # Calculate global usage counts - global_usage = get_global_usage_counts(session, template_ids) - - return [TemplateResponse( - template_id=template.template_id, - template_name=template.template_name, - display_name=template.display_name, - description=template.description, - prompt_template=template.prompt_template, - template_category=template.category, - icon_url=template.icon_url, - is_premium_only=template.is_premium_only, - is_active=template.is_active, - usage_count=global_usage.get(template.template_id, 0), # Global usage count (shows how many times this template has been used overall by users) - created_at=template.created_at, - updated_at=template.updated_at - ) for template in templates] - - finally: - session.close() - - except Exception as e: - logger.log_message(f"Error retrieving templates by category: {str(e)}", level=logging.ERROR) - raise HTTPException(status_code=500, detail=f"Failed to retrieve templates by category: {str(e)}") \ No newline at end of file diff --git a/src/schemas/chat_schema.py b/src/schemas/chat_schemas.py similarity index 60% rename from src/schemas/chat_schema.py rename to src/schemas/chat_schemas.py index 7bbf5615072e4a555fd2faa17ed4d48b3b9f183b..773523a3ea16da92e024730fe74dcb909fa5da5f 100644 --- a/src/schemas/chat_schema.py +++ b/src/schemas/chat_schemas.py @@ -44,25 +44,4 @@ class ChatMessageCreate(BaseModel): message: str model_name: Optional[str] = None system_prompt: Optional[str] = None - formatted_prompt: Optional[str] = None - -class MessageFeedbackCreate(BaseModel): - """Model for creating feedback for a message""" - message_id: int - rating: int # 1-5 star rating - model_name: Optional[str] = None - model_provider: Optional[str] = None - temperature: Optional[float] = None - max_tokens: Optional[int] = None - -class MessageFeedbackResponse(BaseModel): - """Model for feedback response""" - feedback_id: int - message_id: int - rating: Optional[int] = None - model_name: Optional[str] = None - model_provider: Optional[str] = None - temperature: Optional[float] = None - max_tokens: Optional[int] = None - created_at: str - updated_at: str \ No newline at end of file + formatted_prompt: Optional[str] = None \ No newline at end of file diff --git a/src/schemas/code_schema.py b/src/schemas/code_schema.py deleted file mode 100644 index 8a8a51fe749d6ecce1e811c23059c1956bc1b5c8..0000000000000000000000000000000000000000 --- a/src/schemas/code_schema.py +++ /dev/null @@ -1,23 +0,0 @@ -from pydantic import BaseModel -from typing import Optional - -# Request body model -class CodeExecuteRequest(BaseModel): - code: str - session_id: Optional[str] = None - message_id: Optional[int] = None - -class CodeEditRequest(BaseModel): - original_code: str - user_prompt: str - -class CodeFixRequest(BaseModel): - code: str - error: str - -class CodeCleanRequest(BaseModel): - code: str - -class GetLatestCodeRequest(BaseModel): - message_id: int - \ No newline at end of file diff --git a/src/schemas/deep_analysis_schema.py b/src/schemas/deep_analysis_schema.py deleted file mode 100644 index 7a0d9c6395e4749b2a8702af9f95b3c812df9348..0000000000000000000000000000000000000000 --- a/src/schemas/deep_analysis_schema.py +++ /dev/null @@ -1,61 +0,0 @@ -from pydantic import BaseModel -from typing import List, Optional, Any -from datetime import datetime - -# Pydantic models -class DeepAnalysisReportCreate(BaseModel): - report_uuid: str - user_id: Optional[int] = None - goal: str - status: str = "completed" - deep_questions: Optional[str] = None - deep_plan: Optional[str] = None - summaries: Optional[List[Any]] = None - analysis_code: Optional[str] = None - plotly_figures: Optional[List[Any]] = None - synthesis: Optional[List[Any]] = None - final_conclusion: Optional[str] = None - html_report: Optional[str] = None - report_summary: Optional[str] = None - progress_percentage: Optional[int] = 100 - duration_seconds: Optional[int] = None - # Credit and error tracking - credits_consumed: Optional[int] = 0 - error_message: Optional[str] = None - model_provider: Optional[str] = "anthropic" - model_name: Optional[str] = "claude-sonnet-4-6" - total_tokens_used: Optional[int] = 0 - estimated_cost: Optional[float] = 0.0 - steps_completed: Optional[List[str]] = None # Array of completed step names - -class DeepAnalysisReportResponse(BaseModel): - report_id: int - report_uuid: str - user_id: Optional[int] - goal: str - status: str - start_time: datetime - end_time: Optional[datetime] - duration_seconds: Optional[int] - report_summary: Optional[str] - created_at: datetime - updated_at: datetime - -class DeepAnalysisReportDetailResponse(DeepAnalysisReportResponse): - deep_questions: Optional[str] - deep_plan: Optional[str] - summaries: Optional[List[Any]] - analysis_code: Optional[str] - plotly_figures: Optional[List[Any]] - synthesis: Optional[List[Any]] - final_conclusion: Optional[str] - html_report: Optional[str] - progress_percentage: Optional[int] - # Credit and error tracking - credits_consumed: Optional[int] - error_message: Optional[str] - model_provider: Optional[str] = "anthropic" - model_name: Optional[str] = "claude-sonnet-4-6" - total_tokens_used: Optional[int] - estimated_cost: Optional[float] - steps_completed: Optional[List[str]] = None diff --git a/src/schemas/model_settings_schema.py b/src/schemas/model_settings.py similarity index 68% rename from src/schemas/model_settings_schema.py rename to src/schemas/model_settings.py index b0af01bf3d27ea89d4609fe71e160b0ed6435600..be9d97bb335b3ed0b484a064f3fcb9f3cc8d4b3f 100644 --- a/src/schemas/model_settings_schema.py +++ b/src/schemas/model_settings.py @@ -3,5 +3,5 @@ class ModelSettings(BaseModel): provider: str model: str api_key: str = "" - temperature: float = 1.0 - max_tokens: int = 6000 \ No newline at end of file + temperature: float = 0 + max_tokens: int = 1000 \ No newline at end of file diff --git a/src/schemas/query_schema.py b/src/schemas/query_schemas.py similarity index 100% rename from src/schemas/query_schema.py rename to src/schemas/query_schemas.py diff --git a/src/schemas/template_schema.py b/src/schemas/template_schema.py deleted file mode 100644 index 2fdd6c932492468c424e15a1f6523e0a1e487552..0000000000000000000000000000000000000000 --- a/src/schemas/template_schema.py +++ /dev/null @@ -1,36 +0,0 @@ -from pydantic import BaseModel -from typing import Optional -from datetime import datetime - -# Pydantic models for request/response -class TemplateResponse(BaseModel): - template_id: int - template_name: str - display_name: Optional[str] - description: str - prompt_template: str - template_category: Optional[str] - icon_url: Optional[str] - is_premium_only: bool - is_active: bool - usage_count: int - created_at: datetime - updated_at: datetime - -class UserTemplatePreferenceResponse(BaseModel): - template_id: int - template_name: str - display_name: Optional[str] - description: str - template_category: Optional[str] - icon_url: Optional[str] - is_premium_only: bool - is_active: bool - is_enabled: bool - usage_count: int - last_used_at: Optional[datetime] - created_at: Optional[datetime] - updated_at: Optional[datetime] - -class TogglePreferenceRequest(BaseModel): - is_enabled: bool \ No newline at end of file diff --git a/src/schemas/user_schema.py b/src/schemas/user_schemas.py similarity index 100% rename from src/schemas/user_schema.py rename to src/schemas/user_schemas.py diff --git a/src/utils/agent_string_parsing.py b/src/utils/agent_string_parsing.py new file mode 100644 index 0000000000000000000000000000000000000000..2eecb6f45e1452d235fd4455e8aefbb135c6d771 --- /dev/null +++ b/src/utils/agent_string_parsing.py @@ -0,0 +1,30 @@ +import re + +def parse_agents(agent_string): + """ + Parse a string containing agent names separated by ->, (, ), or commas + and return a list of agent names. + """ + if not agent_string or not agent_string.strip(): + return [] + + # Replace parentheses with spaces to handle cases with parentheses + import re + cleaned_string = re.sub(r'\(.*?\)', '', agent_string) + + # Split by -> to get individual agent segments + agent_segments = cleaned_string.split('->') + + # Process each segment to extract agent names + agents = [] + for segment in agent_segments: + # Split by comma and strip whitespace + segment_agents = [agent.strip() for agent in segment.split(',') if agent.strip()] + agents.extend(segment_agents) + + return agents[0] if isinstance(agents, list) else agents + +# sample = """preprocessing_agent(dataset, goal -> code, summary +# instructions='Given a user-defined analysis goal and a pre-loaded dataset df, \nI will generate Python code using NumPy and Pandas to build an exploratory analytics pipeline.\nThe goal is to simplify the preprocessing and introductory analysis of the dataset.\n\nIMPORTANT: You may be provided with previous interaction history. The section marked "##""" + +# print(parse_agents(sample)) \ No newline at end of file diff --git a/src/utils/data/actual-posts.json b/src/utils/data/actual-posts.json deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/src/utils/data/sample-posts.json b/src/utils/data/sample-posts.json deleted file mode 100644 index 77f975c58477c173e2c7ff228581bd83fedc0480..0000000000000000000000000000000000000000 --- a/src/utils/data/sample-posts.json +++ /dev/null @@ -1,38 +0,0 @@ -[ - { - "id": "1", - "title": "Getting Started with Auto-Analyst: A Complete Guide", - "excerpt": "Learn how to set up and use Auto-Analyst for your data analysis needs. This comprehensive guide covers everything from installation to advanced features.", - "content": "# Getting Started with Auto-Analyst: A Complete Guide\n\nWelcome to Auto-Analyst! This comprehensive guide will help you get started with our AI-powered analytics platform.\n\n## What is Auto-Analyst?\n\nAuto-Analyst is an AI-powered analytics platform that helps you analyze data, create visualizations, and generate insights automatically. Our platform combines the power of multiple AI models to provide comprehensive data analysis capabilities.\n\n## Key Features\n\n- **AI-Powered Analysis**: Multiple AI models including GPT-4, Claude, and more\n- **Interactive Visualizations**: Create beautiful charts and graphs\n- **Natural Language Queries**: Ask questions in plain English\n- **Multiple Data Sources**: Connect to various data sources\n- **Collaborative Features**: Share insights with your team\n\n## Getting Started\n\n### 1. Sign Up\nFirst, create your account at [Auto-Analyst](https://auto-analyst.com).\n\n### 2. Upload Your Data\nYou can upload data in various formats:\n- CSV files\n- Excel spreadsheets\n- JSON data\n- Direct database connections\n\n### 3. Start Analyzing\nOnce your data is uploaded, you can start asking questions:\n\n```python\n# Example: Analyze sales data\n\"What are the top-selling products this quarter?\"\n```\n\n## Advanced Features\n\n### Custom Agents\nCreate custom analysis agents for specific use cases:\n\n![Auto-Analyst Dashboard](https://via.placeholder.com/600x300/FF7F7F/FFFFFF?text=Auto-Analyst+Dashboard)\n\n### Deep Analysis\nUse our deep analysis feature for complex insights:\n\n\n\n## Best Practices\n\n1. **Start Simple**: Begin with basic questions and gradually increase complexity\n2. **Use Clear Language**: Be specific in your queries\n3. **Validate Results**: Always verify AI-generated insights\n4. **Iterate**: Refine your analysis based on initial results\n\n## Support\n\nNeed help? Check out our:\n- [Documentation](https://docs.auto-analyst.com)\n- [Community Forum](https://community.auto-analyst.com)\n- [Contact Support](mailto:support@auto-analyst.com)\n\nHappy analyzing! 🚀", - "author": "Auto-Analyst Team", - "publishedAt": "2024-01-15", - "tags": ["tutorial", "getting-started", "guide"], - "featured": true, - "readTime": "8 min read" - }, - { - "id": "2", - "title": "Advanced Data Visualization Techniques", - "excerpt": "Discover advanced techniques for creating compelling data visualizations using Auto-Analyst's powerful visualization tools.", - "content": "# Advanced Data Visualization Techniques\n\nData visualization is a crucial part of data analysis. In this post, we'll explore advanced techniques for creating compelling visualizations.\n\n## Choosing the Right Chart Type\n\nDifferent data types require different visualization approaches:\n\n### Time Series Data\n- Line charts for trends\n- Area charts for cumulative data\n- Candlestick charts for financial data\n\n### Categorical Data\n- Bar charts for comparisons\n- Pie charts for proportions\n- Treemaps for hierarchical data\n\n### Correlation Analysis\n- Scatter plots for relationships\n- Heatmaps for correlation matrices\n- Bubble charts for multi-dimensional data\n\n## Color Theory in Data Visualization\n\nColors play a crucial role in data visualization:\n\n### Color Palettes\n- **Sequential**: For ordered data (light to dark)\n- **Diverging**: For data with a meaningful center\n- **Qualitative**: For categorical data\n\n### Accessibility\nEnsure your visualizations are accessible:\n- Use colorblind-friendly palettes\n- Provide alternative text descriptions\n- Maintain sufficient contrast ratios\n\n## Interactive Features\n\nMake your visualizations interactive:\n\n```javascript\n// Example: Interactive chart configuration\n{\n \"type\": \"scatter\",\n \"data\": {\n \"x\": [1, 2, 3, 4, 5],\n \"y\": [2, 4, 6, 8, 10]\n },\n \"options\": {\n \"interactive\": true,\n \"zoom\": true,\n \"pan\": true\n }\n}\n```\n\n## Best Practices\n\n1. **Keep it Simple**: Don't overcomplicate your visualizations\n2. **Tell a Story**: Use visualizations to support your narrative\n3. **Test with Users**: Get feedback on your visualizations\n4. **Mobile Responsive**: Ensure visualizations work on all devices\n\n## Conclusion\n\nEffective data visualization requires both technical skill and design sensibility. By following these techniques, you can create visualizations that effectively communicate your data insights.", - "author": "Sarah Johnson", - "publishedAt": "2024-01-10", - "tags": ["visualization", "design", "best-practices"], - "featured": false, - "readTime": "6 min read" - }, - { - "id": "3", - "title": "AI Model Comparison: Which One Should You Use?", - "excerpt": "Compare different AI models available in Auto-Analyst and learn when to use each one for optimal results.", - "content": "# AI Model Comparison: Which One Should You Use?\n\nAuto-Analyst supports multiple AI models, each with its own strengths and use cases. Let's compare them:\n\n## Available Models\n\n### GPT-4\n- **Best for**: General analysis, creative tasks\n- **Strengths**: Excellent reasoning, creative problem-solving\n- **Use cases**: Exploratory analysis, hypothesis generation\n\n### Claude (Anthropic)\n- **Best for**: Detailed analysis, code generation\n- **Strengths**: Strong reasoning, helpful responses\n- **Use cases**: Technical analysis, documentation\n\n### Gemini (Google)\n- **Best for**: Multimodal analysis, large datasets\n- **Strengths**: Fast processing, good with images\n- **Use cases**: Quick insights, image analysis\n\n### Groq\n- **Best for**: High-speed processing\n- **Strengths**: Very fast inference\n- **Use cases**: Real-time analysis, rapid prototyping\n\n## Performance Comparison\n\n| Model | Speed | Accuracy | Cost | Best Use Case |\n|-------|-------|----------|------|---------------|\n| GPT-4 | Medium | High | High | General analysis |\n| Claude | Medium | High | Medium | Technical tasks |\n| Gemini | Fast | Medium | Low | Quick insights |\n| Groq | Very Fast | Medium | Low | Real-time analysis |\n\n## Choosing the Right Model\n\n### For Beginners\nStart with **Gemini** or **Groq** for fast, cost-effective analysis.\n\n### For Complex Analysis\nUse **GPT-4** or **Claude** for detailed, nuanced insights.\n\n### For Real-time Applications\nChoose **Groq** for high-speed processing.\n\n## Tips for Model Selection\n\n1. **Start Simple**: Begin with faster, cheaper models\n2. **Iterate**: Try different models for the same task\n3. **Compare Results**: Evaluate outputs from multiple models\n4. **Consider Cost**: Balance performance with budget\n\n## Conclusion\n\nThe best model depends on your specific needs. Experiment with different models to find what works best for your use case.", - "author": "Mike Chen", - "publishedAt": "2024-01-05", - "tags": ["ai", "models", "comparison"], - "featured": false, - "readTime": "5 min read" - } -] - - - diff --git a/src/utils/dataset_description_generator.py b/src/utils/dataset_description_generator.py deleted file mode 100644 index d3df0667ef0e9171b2b185a4dad5f2e947dbfaa0..0000000000000000000000000000000000000000 --- a/src/utils/dataset_description_generator.py +++ /dev/null @@ -1,79 +0,0 @@ -import logging -import pandas as pd -from typing import Dict - -from src.agents.agents import dataset_description_agent, data_context_gen -from src.utils.model_registry import mid_lm -from src.utils.logger import Logger -import dspy - -# Initialize logger -logger = Logger("dataset_description_generator", see_time=False, console_log=False) - -def generate_dataset_description(datasets: Dict[str, pd.DataFrame], existing_description: str = "", dataset_names: list = None) -> str: - """ - Generate AI-powered description for datasets - - Args: - datasets: Dictionary of dataset names to DataFrames - existing_description: Existing description to improve upon (optional) - dataset_names: List of dataset names to use in the description format (optional) - - Returns: - Generated description string with proper exact_python_name formatting - """ - try: - if not datasets or len(datasets) == 0: - return existing_description - - # Build dataset view for description generation - dataset_view = "" - count = 0 - for table_name, table_df in datasets.items(): - head_data = table_df.head(3) - columns = [{col: str(head_data[col].dtype)} for col in head_data.columns] - dataset_view += f"exact_table_name={table_name}\n:columns:{str(columns)}\n{head_data.to_markdown()}\n" - count += 1 - - # Generate description using AI - with dspy.context(lm=mid_lm): - if count == 1: - data_context = dspy.Predict(dataset_description_agent)( - existing_description=existing_description, - dataset=dataset_view - ) - generated_desc = data_context.description - elif count > 1: - data_context = dspy.Predict(data_context_gen)( - user_description=existing_description, - dataset_view=dataset_view - ) - generated_desc = data_context.data_context - else: - generated_desc = existing_description - - # Format the description with exact_python_name for all datasets - if dataset_names and len(dataset_names) > 0: - if len(dataset_names) == 1: - # Single dataset format - formatted_desc = f" exact_python_name: `{dataset_names[0]}` Dataset: {generated_desc}" - else: - # Multiple datasets format - list all dataset names - names_list = ", ".join([f"`{name}`" for name in dataset_names]) - formatted_desc = f" exact_python_name: {names_list} Dataset: {generated_desc}" - else: - # Fallback to original format if no dataset names provided - dataset_keys = list(datasets.keys()) - if len(dataset_keys) == 1: - formatted_desc = f" exact_python_name: `{dataset_keys[0]}` Dataset: {generated_desc}" - else: - names_list = ", ".join([f"`{name}`" for name in dataset_keys]) - formatted_desc = f" exact_python_name: {names_list} Dataset: {generated_desc}" - - logger.log_message(f"Successfully generated dataset description for {count} dataset(s)", level=logging.INFO) - return formatted_desc - - except Exception as e: - logger.log_message(f"Failed to generate dataset description: {str(e)}", level=logging.WARNING) - # Return existing description if generation fails - return existing_description diff --git a/src/utils/generate_report.py b/src/utils/generate_report.py deleted file mode 100644 index 76b944ff7ff23bfdc72e551783342ab2adf30d65..0000000000000000000000000000000000000000 --- a/src/utils/generate_report.py +++ /dev/null @@ -1,820 +0,0 @@ -from bs4 import BeautifulSoup -import markdown -import numpy as np -import re -import pandas as pd - - -def generate_html_report(return_dict): - """Generate a clean HTML report focusing on visualizations and key insights""" - - def convert_markdown_to_html(text): - """Convert markdown text to HTML safely""" - if not text: - return "" - # Don't escape HTML characters before markdown conversion - html = markdown.markdown(str(text), extensions=['tables', 'fenced_code', 'nl2br']) - # Use BeautifulSoup to clean up but preserve structure - soup = BeautifulSoup(html, 'html.parser') - return str(soup) - - def convert_conclusion_to_html(text): - """Special conversion for conclusion with custom bullet point handling""" - if not text: - return "" - - # Clean and prepare text - text = str(text).strip() - - text = re.sub(r'\*\*(.*?)\*\*', r'\1', text) - text = re.sub(r'\*(.*?)\*', r'\1', text) - - # Handle bullet points that might not be properly formatted - lines = text.split('\n') - processed_lines = [] - in_list = False - - for line in lines: - line = line.strip() - if not line: - if in_list: - processed_lines.append('') - in_list = False - processed_lines.append('') - continue - - # Check if line looks like a bullet point - if (line.startswith('- ') or line.startswith('• ') or - line.startswith('* ') or re.match(r'^\d+\.\s', line)): - - if not in_list: - processed_lines.append('
    ') - in_list = True - - # Clean the bullet point - clean_line = re.sub(r'^[-•*]\s*', '', line) - clean_line = re.sub(r'^\d+\.\s*', '', clean_line) - processed_lines.append(f'
  • {clean_line}
    • ') - else: - if in_list: - processed_lines.append('
') - in_list = False - processed_lines.append(f'

{line}

') - - if in_list: - processed_lines.append('') - - # Join and clean up - html_content = '\n'.join(processed_lines) - - # Clean up extra tags and escape HTML entities, but preserve our intentional HTML - html_content = html_content.replace('&', '&').replace('<', '<').replace('>', '>') - - # Restore our intentional HTML tags - html_content = html_content.replace('<strong>', '').replace('</strong>', '') - html_content = html_content.replace('<em>', '').replace('</em>', '') - html_content = html_content.replace('<ul>', '
    ').replace('</ul>', '
') - html_content = html_content.replace('<li>', '
  • ').replace('</li>', '
    • ') - html_content = html_content.replace('<p>', '

    ').replace('</p>', '

    ') - - return html_content - - # Convert key text sections to HTML - goal = convert_markdown_to_html(return_dict['goal']) - questions = convert_markdown_to_html(return_dict['deep_questions']) - conclusion = convert_conclusion_to_html(return_dict['final_conclusion']) - # Remove duplicate conclusion headings and clean up - conclusion = re.sub(r'

    \s*\*\*\s*Conclusion\s*\*\*\s*

    ', '', conclusion, flags=re.IGNORECASE) - conclusion = re.sub(r'\s*Conclusion\s*', '', conclusion, flags=re.IGNORECASE) - conclusion = re.sub(r']*>\s*Conclusion\s*', '', conclusion, flags=re.IGNORECASE) - conclusion = re.sub(r'^\s*Conclusion\s*$', '', conclusion, flags=re.MULTILINE) - - # Combine synthesis content - synthesis_content = '' - if return_dict.get('synthesis'): - synthesis_content = ''.join(f'
    {convert_markdown_to_html(s)}
    ' - for s in return_dict['synthesis']) - - # Generate all visualizations for synthesis section - all_visualizations = [] - if return_dict['plotly_figs']: - for fig_group in return_dict['plotly_figs']: - try: - if isinstance(fig_group, list): - # Handle list of figures - for fig in fig_group: - if hasattr(fig, 'to_html'): - # It's a Plotly Figure object - all_visualizations.append(fig.to_html( - full_html=False, - include_plotlyjs='cdn', - config={'displayModeBar': True} - )) - elif isinstance(fig, str): - # It might be JSON format - try to convert - try: - import plotly.io - fig_obj = plotly.io.from_json(fig) - all_visualizations.append(fig_obj.to_html( - full_html=False, - include_plotlyjs='cdn', - config={'displayModeBar': True} - )) - except Exception as e: - print(f"Warning: Could not process figure JSON: {e}") - continue - else: - # Single figure - if hasattr(fig_group, 'to_html'): - # It's a Plotly Figure object - all_visualizations.append(fig_group.to_html( - full_html=False, - include_plotlyjs='cdn', - config={'displayModeBar': True} - )) - elif isinstance(fig_group, str): - # It might be JSON format - try to convert - try: - import plotly.io - fig_obj = plotly.io.from_json(fig_group) - all_visualizations.append(fig_obj.to_html( - full_html=False, - include_plotlyjs='cdn', - config={'displayModeBar': True} - )) - except Exception as e: - print(f"Warning: Could not process figure JSON: {e}") - continue - - except Exception as e: - print(f"Warning: Error processing visualizations: {e}") - - # Prepare code for syntax highlighting - code_content = return_dict.get('code', '').strip() - - html = f""" - - - - Deep Analysis Report - - - - - - - - - - -
    -
    -

    Deep Analysis Report

    -

    Original Question

    -
    - {goal} -
    - -

    Detailed Research Questions

    -
    - {questions} -
    -
    - - -
    -

    Analysis & Insights

    -
    - {synthesis_content} -
    - - {''.join(f'
    {viz}
    ' for viz in all_visualizations) if all_visualizations else '

    No visualizations generated

    '} -
    - - {f''' -
    -

    Generated Code

    -
    -
    - - View Generated Code (Click to expand) - -
    - -
    -
    -
    - 
    {code_content}
    -
    -
    -
    - ''' if code_content else ''} - -
    -

    Conclusion

    -
    - {conclusion} -
    -
    -
    - - -
    - - """ - return html - diff --git a/src/utils/logger.py b/src/utils/logger.py index eb877c73bd1c3d606fd054d15d77a44e068c84aa..747dced2f46994c973df322e04f2f90915ae53e6 100644 --- a/src/utils/logger.py +++ b/src/utils/logger.py @@ -1,14 +1,13 @@ import os import time import logging -import sys from dotenv import load_dotenv load_dotenv() class Logger: def __init__(self, name: str, see_time: bool = False, console_log: bool = False, level: int = logging.INFO): - self.is_dev = os.getenv("ENVIRONMENT", "development") == "development" + self.is_dev = os.getenv("ENV", "production") == "development" self.logger = logging.getLogger(name) self.logger.setLevel(level) @@ -22,47 +21,28 @@ class Logger: "%(asctime)s - %(name)s - %(levelname)s - %(message)s" if see_time else "%(message)s," ) - # File handler with UTF-8 encoding - file_handler = logging.FileHandler(f"./logs/{name}.log", encoding='utf-8') + file_handler = logging.FileHandler(f"./logs/{name}.log") file_handler.setFormatter(formatter) self.logger.addHandler(file_handler) if console_log: - # Console handler with UTF-8 encoding - console_handler = logging.StreamHandler(sys.stdout) + console_handler = logging.StreamHandler() console_handler.setFormatter(formatter) - # Force UTF-8 encoding for console output - if hasattr(console_handler.stream, 'reconfigure'): - console_handler.stream.reconfigure(encoding='utf-8') self.logger.addHandler(console_handler) def log_message(self, message: str, level: int = logging.INFO): if not self.is_dev: return - try: - if level == logging.INFO: - self.logger.info(message) - elif level == logging.ERROR: - self.logger.error(message) - elif level == logging.WARNING: - self.logger.warning(message) - elif level == logging.DEBUG: - self.logger.debug(message) - else: - self.logger.info(message) - except UnicodeEncodeError: - # Fallback: remove emoji characters if encoding fails - safe_message = message.encode('ascii', 'ignore').decode('ascii') - if level == logging.INFO: - self.logger.info(safe_message) - elif level == logging.ERROR: - self.logger.error(safe_message) - elif level == logging.WARNING: - self.logger.warning(safe_message) - elif level == logging.DEBUG: - self.logger.debug(safe_message) - else: - self.logger.info(safe_message) + if level == logging.INFO: + self.logger.info(message) + elif level == logging.ERROR: + self.logger.error(message) + elif level == logging.WARNING: + self.logger.warning(message) + elif level == logging.DEBUG: + self.logger.debug(message) + else: + self.logger.info(message) def disable_logging(self): self.logger.disabled = True diff --git a/src/utils/model_registry.py b/src/utils/model_registry.py index bc61bee315a36ab151235695efa4224369a9a3a3..42c839587207995e92555450fa948bf6511731f7 100644 --- a/src/utils/model_registry.py +++ b/src/utils/model_registry.py @@ -1,5 +1,7 @@ -import dspy -import os +""" +Models registry for the Auto-Analyst application. +This file serves as the single source of truth for all model information. +""" # Model providers PROVIDERS = { @@ -8,338 +10,145 @@ PROVIDERS = { "groq": "GROQ", "gemini": "Google Gemini" } -max_tokens = int(os.getenv("MAX_TOKENS", 6000)) -# Clamp temperature to valid range (0..1) for all models -default_temperature = min(1.0, max(0.0, float(os.getenv("TEMPERATURE", "1.0")))) - -# OpenAI reasoning models (gpt-5 family, o3, etc.) only accept temperature=1.0 (or None). -# dspy>=3.2 validates this at dspy.LM(...) construction, so never pass the env-derived -# temperature to these models or the app will fail to import when TEMPERATURE != 1.0. -reasoning_temperature = 1.0 - -# Lightweight LMs used for small internal tasks (planning, classification, etc.) -small_lm = dspy.LM('anthropic/claude-haiku-4-6', temperature=default_temperature, max_tokens=300, api_key=os.getenv("ANTHROPIC_API_KEY"), cache=False) - -mid_lm = dspy.LM('anthropic/claude-haiku-4-6', temperature=default_temperature, max_tokens=1800, api_key=os.getenv("ANTHROPIC_API_KEY"), cache=False) - -# OpenAI models -gpt_5_nano = dspy.LM( - model="openai/gpt-5-nano", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=reasoning_temperature, - max_tokens=16_000, - cache=False -) - -gpt_5_mini = dspy.LM( - model="openai/gpt-5-mini", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=reasoning_temperature, - max_tokens=16_000, - cache=False -) - -gpt_5 = dspy.LM( - model="openai/gpt-5", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=reasoning_temperature, - max_tokens=16_000, - cache=False -) - -gpt_5_2 = dspy.LM( - model="openai/gpt-5.2", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=reasoning_temperature, - max_tokens=max(max_tokens, 16000), - cache=False -) - -gpt_5_2_pro = dspy.LM( - model="openai/gpt-5.2-pro", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=reasoning_temperature, - max_tokens=max(max_tokens, 16000), - cache=False -) - -gpt_5_2_chat_latest = dspy.LM( - model="openai/gpt-5.2-chat-latest", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=default_temperature, - max_tokens=max(max_tokens, 16000), - cache=False -) - -gpt_5_4 = dspy.LM( - model="openai/gpt-5.4", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=reasoning_temperature, - max_tokens=16_000, - cache=False -) - -gpt_5_4_pro = dspy.LM( - model="openai/gpt-5.4-pro", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=reasoning_temperature, - max_tokens=16_000, - cache=False -) - -o3 = dspy.LM( - model="openai/o3-2025-04-16", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=reasoning_temperature, - max_tokens=20_000, - cache=False -) - -# Anthropic models -claude_haiku_4_5 = dspy.LM( - model="anthropic/claude-haiku-4-5-20251001", - api_key=os.getenv("ANTHROPIC_API_KEY"), - temperature=default_temperature, - max_tokens=max_tokens, - cache=False -) - -claude_sonnet_4_5 = dspy.LM( - model="anthropic/claude-sonnet-4-5-20250929", - api_key=os.getenv("ANTHROPIC_API_KEY"), - temperature=default_temperature, - max_tokens=max_tokens, - cache=False -) - -claude_sonnet_4_6 = dspy.LM( - model="anthropic/claude-sonnet-4-6", - api_key=os.getenv("ANTHROPIC_API_KEY"), - temperature=default_temperature, - max_tokens=max_tokens, - cache=False -) - -claude_opus_4_5 = dspy.LM( - model="anthropic/claude-opus-4-5-20251101", - api_key=os.getenv("ANTHROPIC_API_KEY"), - temperature=float(os.getenv("TEMPERATURE", 1.0)), - max_tokens=max_tokens, - cache=False -) - -claude_opus_4_6 = dspy.LM( - model="anthropic/claude-opus-4-6", - api_key=os.getenv("ANTHROPIC_API_KEY"), - temperature=default_temperature, - max_tokens=max_tokens, - cache=False -) - -# Groq models -deepseek_r1_distill_llama_70b = dspy.LM( - model="groq/deepseek-r1-distill-llama-70b", - api_key=os.getenv("GROQ_API_KEY"), - temperature=default_temperature, - max_tokens=max_tokens, - cache=False -) - -gpt_oss_120B = dspy.LM( - model="groq/gpt-oss-120B", - api_key=os.getenv("GROQ_API_KEY"), - temperature=default_temperature, - max_tokens=max_tokens, - cache=False -) - -gpt_oss_20B = dspy.LM( - model="groq/gpt-oss-20B", - api_key=os.getenv("GROQ_API_KEY"), - temperature=default_temperature, - max_tokens=max_tokens, - cache=False -) - -# Gemini models -gemini_2_5_pro_preview_03_25 = dspy.LM( - model="gemini/gemini-2.5-pro-preview-03-25", - api_key=os.getenv("GEMINI_API_KEY"), - temperature=default_temperature, - max_tokens=max_tokens, - cache=False -) - -gemini_3_pro = dspy.LM( - model="gemini/gemini-3-pro", - api_key=os.getenv("GEMINI_API_KEY"), - temperature=float(os.getenv("TEMPERATURE", 1.0)), - max_tokens=max_tokens, - cache=False -) - -gemini_3_flash = dspy.LM( - model="gemini/gemini-3-flash", - api_key=os.getenv("GEMINI_API_KEY"), - temperature=float(os.getenv("TEMPERATURE", 1.0)), - max_tokens=max_tokens, - cache=False -) - -MODEL_OBJECTS = { - # OpenAI models - "gpt-5-nano": gpt_5_nano, - "gpt-5-mini": gpt_5_mini, - "gpt-5": gpt_5, - "gpt-5.2": gpt_5_2, - "gpt-5.2-pro": gpt_5_2_pro, - "gpt-5.2-chat-latest": gpt_5_2_chat_latest, - "gpt-5.4": gpt_5_4, - "gpt-5.4-pro": gpt_5_4_pro, - "o3": o3, - - # Anthropic models - "claude-haiku-4-5": claude_haiku_4_5, - "claude-sonnet-4-5-20250929": claude_sonnet_4_5, - "claude-sonnet-4-6": claude_sonnet_4_6, - "claude-opus-4-5-20251101": claude_opus_4_5, - "claude-opus-4-6": claude_opus_4_6, - - # Groq models - "deepseek-r1-distill-llama-70b": deepseek_r1_distill_llama_70b, - "gpt-oss-120B": gpt_oss_120B, - "gpt-oss-20B": gpt_oss_20B, - - # Gemini models - "gemini-2.5-pro-preview-03-25": gemini_2_5_pro_preview_03_25, - "gemini-3-pro": gemini_3_pro, - "gemini-3-flash": gemini_3_flash +# Cost per 1K tokens for different models +MODEL_COSTS = { + "openai": { + "gpt-4.1": {"input": 0.002, "output": 0.008}, + "gpt-4.1-mini": {"input": 0.0004, "output": 0.0016}, + "gpt-4.1-nano": {"input": 0.00010, "output": 0.0004}, + "gpt-4.5-preview": {"input": 0.075, "output": 0.15}, + "gpt-4o": {"input": 0.0025, "output": 0.01}, + "gpt-4o-mini": {"input": 0.00015, "output": 0.0006}, + "o1": {"input": 0.015, "output": 0.06}, + "o1-pro": {"input": 0.015, "output": 0.6}, + "o1-mini": {"input": 0.00011, "output": 0.00044}, + "o3": {"input": 0.001, "output": 0.04}, + "o3-mini": {"input": 0.00011, "output": 0.00044}, + "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015}, + }, + "anthropic": { + "claude-3-opus-latest": {"input": 0.015, "output": 0.075}, + "claude-3-7-sonnet-latest": {"input": 0.003, "output": 0.015}, + "claude-3-5-sonnet-latest": {"input": 0.003, "output": 0.015}, + "claude-3-5-haiku-latest": {"input": 0.0008, "output": 0.0004}, + }, + "groq": { + "deepseek-r1-distill-llama-70b": {"input": 0.00075, "output": 0.00099}, + "llama-3.3-70b-versatile": {"input": 0.00059, "output": 0.00079}, + "llama3-8b-8192": {"input": 0.00005, "output": 0.00008}, + "llama3-70b-8192": {"input": 0.00059, "output": 0.00079}, + "mistral-saba-24b": {"input": 0.00079, "output": 0.00079}, + "gemma2-9b-it": {"input": 0.0002, "output": 0.0002}, + "qwen-qwq-32b": {"input": 0.00029, "output": 0.00039}, + "meta-llama/llama-4-maverick-17b-128e-instruct": {"input": 0.0002, "output": 0.0006}, + "meta-llama/llama-4-scout-17b-16e-instruct": {"input": 0.00011, "output": 0.00034}, + "deepseek-r1-distill-qwen-32b": {"input": 0.00075, "output": 0.00099}, + "llama-3.1-70b-versatile": {"input": 0.00059, "output": 0.00079}, + }, + "gemini": { + "gemini-2.5-pro-preview-03-25": {"input": 0.00015, "output": 0.001} + } } - -def get_model_object(model_name: str): - """Get model object by name""" - return MODEL_OBJECTS.get(model_name, claude_sonnet_4_6) - - -# Get max tokens from environment -max_tokens = int(os.getenv("MAX_TOKENS", 6000)) - # Tiers based on cost per 1K tokens MODEL_TIERS = { "tier1": { "name": "Basic", "credits": 1, "models": [ - "gpt-5-nano", - "gpt-oss-20B" + "llama3-8b-8192", + "llama-3.1-8b-instant", + "gemma2-9b-it", + "meta-llama/llama-4-scout-17b-16e-instruct", + "llama-3.2-1b-preview", + "llama-3.2-3b-preview", + "llama-3.2-11b-text-preview", + "llama-3.2-11b-vision-preview", + "llama3-groq-8b-8192-tool-use-preview" ] }, "tier2": { "name": "Standard", "credits": 3, "models": [ - "claude-haiku-4-5", - "gpt-5-mini", - "gpt-5.2-chat-latest" + "gpt-4.1-nano", + "gpt-4o-mini", + "o1-mini", + "o3-mini", + "qwen-qwq-32b", + "meta-llama/llama-4-maverick-17b-128e-instruct" ] }, "tier3": { "name": "Premium", "credits": 5, "models": [ + "gpt-4.1", + "gpt-4.1-mini", + "gpt-4.5-preview", + "gpt-4o", + "o1", + "o1-pro", "o3", - "claude-sonnet-4-5-20250929", - "claude-sonnet-4-6", + "gpt-3.5-turbo", + "claude-3-opus-latest", + "claude-3-7-sonnet-latest", + "claude-3-5-sonnet-latest", + "claude-3-5-haiku-latest", "deepseek-r1-distill-llama-70b", - "gpt-oss-120B", - "gemini-2.5-pro-preview-03-25", - "gemini-3-flash", - "gpt-5.2" - ] - }, - "tier4": { - "name": "Premium Plus", - "credits": 20, - "models": [ - "gpt-5", - "gpt-5.4", - "claude-opus-4-5-20251101", - "claude-opus-4-6", - "gemini-3-pro" - ] - }, - "tier5": { - "name": "Ultimate", - "credits": 50, - "models": [ - "gpt-5.2-pro", - "gpt-5.4-pro" + "llama-3.3-70b-versatile", + "llama3-70b-8192", + "mistral-saba-24b", + "deepseek-r1-distill-qwen-32b", + "llama-3.2-90b-text-preview", + "llama-3.2-90b-vision-preview", + "llama-3.3-70b-specdec", + "llama2-70b-4096", + "llama-3.1-70b-versatile", + "llama-3.1-405b-reasoning", + "llama3-groq-70b-8192-tool-use-preview", + "gemini-2.5-pro-preview-03-25" ] } } + # Model metadata (display name, context window, etc.) MODEL_METADATA = { # OpenAI - "gpt-5-nano": {"display_name": "GPT-5 Nano", "context_window": 64000}, - "gpt-5-mini": {"display_name": "GPT-5 Mini", "context_window": 150000}, - "gpt-5": {"display_name": "GPT-5", "context_window": 400000}, - "gpt-5.2": {"display_name": "GPT-5.2", "context_window": 400000}, - "gpt-5.2-pro": {"display_name": "GPT-5.2 Pro", "context_window": 400000}, - "gpt-5.2-chat-latest": {"display_name": "GPT-5.2 Chat", "context_window": 400000}, - "gpt-5.4": {"display_name": "GPT-5.4", "context_window": 1050000}, - "gpt-5.4-pro": {"display_name": "GPT-5.4 Pro", "context_window": 1050000}, + "gpt-4.1": {"display_name": "GPT-4.1", "context_window": 128000}, + "gpt-4.1-mini": {"display_name": "GPT-4.1 Mini", "context_window": 128000}, + "gpt-4.1-nano": {"display_name": "GPT-4.1 Nano", "context_window": 128000}, + "gpt-4o": {"display_name": "GPT-4o", "context_window": 128000}, + "gpt-4.5-preview": {"display_name": "GPT-4.5 Preview", "context_window": 128000}, + "gpt-4o-mini": {"display_name": "GPT-4o Mini", "context_window": 128000}, + "gpt-3.5-turbo": {"display_name": "GPT-3.5 Turbo", "context_window": 16385}, + "o1": {"display_name": "o1", "context_window": 128000}, + "o1-pro": {"display_name": "o1 Pro", "context_window": 128000}, + "o1-mini": {"display_name": "o1 Mini", "context_window": 128000}, "o3": {"display_name": "o3", "context_window": 128000}, - + "o3-mini": {"display_name": "o3 Mini", "context_window": 128000}, # Anthropic - "claude-haiku-4-5": {"display_name": "Claude Haiku 4.5", "context_window": 200000}, - "claude-sonnet-4-5-20250929": {"display_name": "Claude Sonnet 4.5", "context_window": 200000}, - "claude-sonnet-4-6": {"display_name": "Claude Sonnet 4.6", "context_window": 1000000}, - "claude-opus-4-5-20251101": {"display_name": "Claude Opus 4.5", "context_window": 200000}, - "claude-opus-4-6": {"display_name": "Claude Opus 4.6", "context_window": 1000000}, - + "claude-3-opus-latest": {"display_name": "Claude 3 Opus", "context_window": 200000}, + "claude-3-7-sonnet-latest": {"display_name": "Claude 3.7 Sonnet", "context_window": 200000}, + "claude-3-5-sonnet-latest": {"display_name": "Claude 3.5 Sonnet", "context_window": 200000}, + "claude-3-5-haiku-latest": {"display_name": "Claude 3.5 Haiku", "context_window": 200000}, + # GROQ "deepseek-r1-distill-llama-70b": {"display_name": "DeepSeek R1 Distill Llama 70b", "context_window": 32768}, - "gpt-oss-120B": {"display_name": "OpenAI gpt oss 120B", "context_window": 128000}, - "gpt-oss-20B": {"display_name": "OpenAI gpt oss 20B", "context_window": 128000}, - + "llama-3.3-70b-versatile": {"display_name": "Llama 3.3 70b", "context_window": 8192}, + "llama3-8b-8192": {"display_name": "Llama 3 8b", "context_window": 8192}, + "llama3-70b-8192": {"display_name": "Llama 3 70b", "context_window": 8192}, + "mistral-saba-24b": {"display_name": "Mistral Saba 24b", "context_window": 32768}, + "gemma2-9b-it": {"display_name": "Gemma 2 9b", "context_window": 8192}, + "qwen-qwq-32b": {"display_name": "Qwen QWQ 32b | Alibaba", "context_window": 32768}, + "meta-llama/llama-4-maverick-17b-128e-instruct": {"display_name": "Llama 4 Maverick 17b", "context_window": 128000}, + "meta-llama/llama-4-scout-17b-16e-instruct": {"display_name": "Llama 4 Scout 17b", "context_window": 16000}, + "llama-3.1-70b-versatile": {"display_name": "Llama 3.1 70b Versatile", "context_window": 8192}, + # Gemini "gemini-2.5-pro-preview-03-25": {"display_name": "Gemini 2.5 Pro", "context_window": 1000000}, - "gemini-3-pro": {"display_name": "Gemini 3 Pro", "context_window": 1000000}, - "gemini-3-flash": {"display_name": "Gemini 3 Flash", "context_window": 1000000}, -} - -MODEL_COSTS = { - "openai": { - "gpt-5-nano": {"input": 0.00005, "output": 0.0004}, - "gpt-5-mini": {"input": 0.00025, "output": 0.002}, - "gpt-5": {"input": 0.00125, "output": 0.01}, - "gpt-5.2": {"input": 0.00125, "output": 0.01}, - "gpt-5.2-pro": {"input": 0.002, "output": 0.015}, - "gpt-5.2-chat-latest": {"input": 0.0005, "output": 0.002}, - "gpt-5.4": {"input": 0.0025, "output": 0.015}, - "gpt-5.4-pro": {"input": 0.03, "output": 0.18}, - "o3": {"input": 0.002, "output": 0.008}, - }, - "anthropic": { - "claude-haiku-4-5": {"input": 0.001, "output": 0.005}, - "claude-sonnet-4-5-20250929": {"input": 0.003, "output": 0.015}, - "claude-sonnet-4-6": {"input": 0.003, "output": 0.015}, - "claude-opus-4-5-20251101": {"input": 0.015, "output": 0.075}, - "claude-opus-4-6": {"input": 0.005, "output": 0.025}, - }, - "groq": { - "deepseek-r1-distill-llama-70b": {"input": 0.00075, "output": 0.00099}, - "gpt-oss-120B": {"input": 0.00075, "output": 0.00099}, - "gpt-oss-20B": {"input": 0.00075, "output": 0.00099} - }, - "gemini": { - "gemini-2.5-pro-preview-03-25": {"input": 0.00015, "output": 0.001}, - "gemini-3-pro": {"input": 0.0002, "output": 0.001}, - "gemini-3-flash": {"input": 0.0001, "output": 0.0005} - } } # Helper functions @@ -400,4 +209,4 @@ def get_all_models_for_provider(provider): def get_models_by_tier(tier_id): """Get all models for a specific tier""" - return MODEL_TIERS.get(tier_id, {}).get("models", []) + return MODEL_TIERS.get(tier_id, {}).get("models", []) \ No newline at end of file diff --git a/src/utils/simple_retriever.py b/src/utils/simple_retriever.py deleted file mode 100644 index 750cd524bf58d6467c0abc2f0167af79c3014734..0000000000000000000000000000000000000000 --- a/src/utils/simple_retriever.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Lightweight TF-IDF based retriever to replace llama-index VectorStoreIndex. -Uses scikit-learn (already a dependency) — no heavyweight RAG framework needed. -""" - -import numpy as np -from sklearn.feature_extraction.text import TfidfVectorizer -from sklearn.metrics.pairwise import cosine_similarity - - -class TextNode: - """Minimal node wrapper matching the llama-index retrieval result interface.""" - - def __init__(self, text: str): - self.text = text - - -class Document: - """Minimal document wrapper matching the llama-index Document interface.""" - - def __init__(self, text: str): - self.text = text - - -class SimpleRetriever: - """ - TF-IDF cosine-similarity retriever that mirrors the subset of the - llama-index VectorStoreIndex / BaseRetriever interface actually used: - - index = SimpleRetriever.from_documents([Document(text=x) for x in docs]) - ret = index.as_retriever(similarity_top_k=1) - nodes = ret.retrieve(query) # -> list[TextNode] - text = nodes[0].text - """ - - def __init__(self, texts: list[str]): - self._docs = texts - self._top_k = 1 - self._vectorizer = TfidfVectorizer() - if texts: - self._matrix = self._vectorizer.fit_transform(texts) - else: - self._matrix = None - - @classmethod - def from_documents(cls, documents) -> "SimpleRetriever": - """Build a retriever from a list of Document objects.""" - texts = [doc.text for doc in documents] - return cls(texts) - - def as_retriever(self, similarity_top_k: int = 1) -> "SimpleRetriever": - """Configure top-k and return self (mirrors llama-index fluent API).""" - self._top_k = similarity_top_k - return self - - def retrieve(self, query: str) -> list[TextNode]: - """Return the top-k most similar nodes for a query.""" - if not self._docs or self._matrix is None: - return [TextNode("")] - - query_vec = self._vectorizer.transform([query]) - scores = cosine_similarity(query_vec, self._matrix)[0] - top_indices = np.argsort(scores)[::-1][: self._top_k] - return [TextNode(self._docs[i]) for i in top_indices] diff --git a/src/utils/test.ipynb b/src/utils/test.ipynb deleted file mode 100644 index 8e8b221249915f4bc9a80483672ee16e660ed07e..0000000000000000000000000000000000000000 --- a/src/utils/test.ipynb +++ /dev/null @@ -1,109 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "id": "fc00c2fa", - "metadata": {}, - "outputs": [], - "source": [ - "import dspy\n", - "import os" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3a1670ba", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null - "id": "194bf39b", - "metadata": {}, - "outputs": [], - "source": [ - "gpt_5_mini = dspy.LM(\n", - " model=\"openai/gpt-5-mini\",\n", - " api_key=os.getenv(\"OPENAI_API_KEY\"),\n", - " temperature=float(os.getenv(\"TEMPERATURE\", 1.0)),\n", - " max_tokens= None,\n", - " # max_completion_tokens=3000,\n", - " cache=False\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "011b0ca4", - "metadata": {}, - "outputs": [], - "source": [ - "dspy.configure(lm= gpt_5_mini)\n", - "\n", - "allocator = dspy.Predict(\"goal,planner_desc,dataset->exact_word_complexity:Literal['unrelated','basic', 'intermediate', 'advanced'],reasoning\")\n", - "\n", - "\n", - "session = allocator(goal=\"build me a regression model and then visualize the residuals\", planner_desc='{data_viz_agent:\"I love visualizaing\"}', dataset='housing dataset, can only answer housing')\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "0217d501", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "Prediction(\n", - " exact_word_complexity='intermediate',\n", - " reasoning='Building a regression model and visualizing residuals requires moderate statistical and coding knowledge: selecting and preprocessing features from the housing dataset, splitting data, fitting a model (e.g., linear or regularized regression), computing evaluation metrics (MSE, R²), and creating diagnostic plots (residuals vs. fitted, histogram or KDE of residuals, Q–Q plot, and scale-location). It’s more than a basic copy-paste task because it involves assumption checks and exploratory/feature work, but it doesn’t require the advanced theory or highly specialized modeling techniques that would push it into “advanced.” The work is therefore best categorized as intermediate.'\n", - ")" - ] - }, - "execution_count": 12, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "session" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2bf84f7e", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "base", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.7" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/test_code_fix.py b/test_code_fix.py new file mode 100644 index 0000000000000000000000000000000000000000..3d81dde5d1360eb6b576983f50ca03f863af9409 --- /dev/null +++ b/test_code_fix.py @@ -0,0 +1,84 @@ +import sys +import os +from src.routes.code_routes import fix_code_with_dspy, extract_code_blocks +from scripts.format_response import execute_code_from_markdown +import pandas as pd +import logging + +# Configure logging +logging.basicConfig(level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + +def read_code_file(file_path): + """Read the code file content""" + with open(file_path, 'r') as f: + return f.read() + +def demo_code_fixing(): + """Demonstrate the code fixing functionality using sample code with errors""" + + print("\n===== CODE ERROR FIXING DEMONSTRATION =====\n") + + # Sample dataset + sample_data = { + 'price': [10000, 12000, 13000, 15000, 18000], + 'area': [1000, 1200, 1300, 1500, 1800], + 'bedrooms': [2, 3, 3, 4, 4], + 'bathrooms': [1, 1, 2, 2, 3], + 'stories': [1, 1, 2, 2, 3], + 'mainroad': ['yes', 'yes', 'no', 'yes', 'yes'], + 'guestroom': ['no', 'no', 'yes', 'no', 'yes'], + 'basement': ['no', 'no', 'yes', 'no', 'yes'], + 'hotwaterheating': ['no', 'yes', 'no', 'yes', 'yes'], + 'airconditioning': ['yes', 'yes', 'no', 'yes', 'yes'], + 'parking': [1, 1, 2, 2, 3], + 'prefarea': ['yes', 'no', 'yes', 'no', 'yes'], + 'furnishingstatus': ['furnished', 'semi-furnished', 'unfurnished', 'furnished', 'semi-furnished'] + } + + df = pd.DataFrame(sample_data) + + # 1. Load the sample code with errors + sample_code_path = "sample_code.py" + original_code = read_code_file(sample_code_path) + + # Introduce an error in the statistical analytics agent + # Change numpy import from np to pd (conflict with pandas) + modified_code = original_code.replace("import numpy as np", "import numpy as pd") + + print("Original code has been modified to introduce an error (numpy import conflict)") + + # 2. Execute the code to get the error + print("\n=== EXECUTING CODE WITH ERROR... ===\n") + output, _ = execute_code_from_markdown(modified_code, df) + print(output) + + # 3. Fix the code using our error fixing function + print("\n=== FIXING THE CODE... ===\n") + dataset_context = f"DataFrame with {len(df)} rows and columns: {', '.join(df.columns)}" + fixed_code = fix_code_with_dspy(modified_code, output, dataset_context) + + # 4. Execute the fixed code to verify it works + print("\n=== EXECUTING FIXED CODE... ===\n") + fixed_output, _ = execute_code_from_markdown(fixed_code, df) + print(fixed_output) + + # 5. Show what changes were made + print("\n=== ANALYZING FIXES MADE ===\n") + original_blocks = extract_code_blocks(modified_code) + fixed_blocks = extract_code_blocks(fixed_code) + + for block_name in original_blocks: + if block_name in fixed_blocks: + print(f"Block: {block_name}") + # Very simple diff - just check if they're different + if original_blocks[block_name] != fixed_blocks[block_name]: + print(f" - This block was modified during fixing") + else: + print(f" - No changes to this block") + + print("\n===== DEMONSTRATION COMPLETE =====") + +if __name__ == "__main__": + demo_code_fixing() \ No newline at end of file diff --git a/test_datasets.ipynb b/test_datasets.ipynb deleted file mode 100644 index 47c9af7fdbb3365636348ab77f075fc07633e36c..0000000000000000000000000000000000000000 --- a/test_datasets.ipynb +++ /dev/null @@ -1,869 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "ae0e83db", - "metadata": {}, - "outputs": [], - "source": [ - "import dspy \n", - "import duckdb \n", - "import pandas as pd\n", - "import numpy as np\n", - "import os\n", - "from dotenv import load_dotenv\n", - "\n", - "load_dotenv()\n", - "\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "94387056", - "metadata": {}, - "outputs": [], - "source": [ - "styling_instructions = [\n", - " {\n", - " \"category\": \"line_charts\",\n", - " \"description\": \"Used to visualize trends and changes over time, often with multiple series.\",\n", - " \"styling\": {\n", - " \"template\": \"plotly_white\",\n", - " \"axes_line_width\": 0.2,\n", - " \"grid_width\": 1,\n", - " \"title\": {\n", - " \"bold_html\": True,\n", - " \"include\": True\n", - " },\n", - " \"colors\": \"use multiple colors if more than one line\",\n", - " \"annotations\": [\"min\", \"max\"],\n", - " \"number_format\": {\n", - " \"apply_k_m\": True,\n", - " \"thresholds\": {\"K\": 1000, \"M\": 100000},\n", - " \"percentage_decimals\": 2,\n", - " \"percentage_sign\": True\n", - " },\n", - " \"default_size\": {\"height\": 1200, \"width\": 1000}\n", - " }\n", - " },\n", - " {\n", - " \"category\": \"bar_charts\",\n", - " \"description\": \"Useful for comparing discrete categories or groups with bars representing values.\",\n", - " \"styling\": {\n", - " \"template\": \"plotly_white\",\n", - " \"axes_line_width\": 0.2,\n", - " \"grid_width\": 1,\n", - " \"title\": {\"bold_html\": True, \"include\": True},\n", - " \"annotations\": [\"bar values\"],\n", - " \"number_format\": {\n", - " \"apply_k_m\": True,\n", - " \"thresholds\": {\"K\": 1000, \"M\": 100000},\n", - " \"percentage_decimals\": 2,\n", - " \"percentage_sign\": True\n", - " },\n", - " \"default_size\": {\"height\": 1200, \"width\": 1000}\n", - " }\n", - " },\n", - " {\n", - " \"category\": \"histograms\",\n", - " \"description\": \"Display the distribution of a data set, useful for returns or frequency distributions.\",\n", - " \"styling\": {\n", - " \"template\": \"plotly_white\",\n", - " \"bin_size\": 50,\n", - " \"axes_line_width\": 0.2,\n", - " \"grid_width\": 1,\n", - " \"title\": {\"bold_html\": True, \"include\": True},\n", - " \"annotations\": [\"x values\"],\n", - " \"number_format\": {\n", - " \"apply_k_m\": True,\n", - " \"thresholds\": {\"K\": 1000, \"M\": 100000},\n", - " \"percentage_decimals\": 2,\n", - " \"percentage_sign\": True\n", - " },\n", - " \"default_size\": {\"height\": 1200, \"width\": 1000}\n", - " }\n", - " },\n", - " {\n", - " \"category\": \"pie_charts\",\n", - " \"description\": \"Show composition or parts of a whole with slices representing categories.\",\n", - " \"styling\": {\n", - " \"template\": \"plotly_white\",\n", - " \"top_categories_to_show\": 10,\n", - " \"bundle_rest_as\": \"Others\",\n", - " \"axes_line_width\": 0.2,\n", - " \"grid_width\": 1,\n", - " \"title\": {\"bold_html\": True, \"include\": True},\n", - " \"annotations\": [\"x values\"],\n", - " \"number_format\": {\n", - " \"apply_k_m\": True,\n", - " \"thresholds\": {\"K\": 1000, \"M\": 100000},\n", - " \"percentage_decimals\": 2,\n", - " \"percentage_sign\": True\n", - " },\n", - " \"default_size\": {\"height\": 1200, \"width\": 1000}\n", - " }\n", - " },\n", - " {\n", - " \"category\": \"tabular_and_generic_charts\",\n", - " \"description\": \"Applies to charts where number formatting needs flexibility, including mixed or raw data.\",\n", - " \"styling\": {\n", - " \"template\": \"plotly_white\",\n", - " \"axes_line_width\": 0.2,\n", - " \"grid_width\": 1,\n", - " \"title\": {\"bold_html\": True, \"include\": True},\n", - " \"annotations\": [\"x values\"],\n", - " \"number_format\": {\n", - " \"apply_k_m\": True,\n", - " \"thresholds\": {\"K\": 1000, \"M\": 100000},\n", - " \"exclude_if_commas_present\": True,\n", - " \"exclude_if_not_numeric\": True,\n", - " \"percentage_decimals\": 2,\n", - " \"percentage_sign\": True\n", - " },\n", - " \"default_size\": {\"height\": 1200, \"width\": 1000}\n", - " }\n", - " },\n", - " {\n", - " \"category\": \"heat_maps\",\n", - " \"description\": \"Show data density or intensity using color scales on a matrix or grid.\",\n", - " \"styling\": {\n", - " \"template\": \"plotly_white\",\n", - " \"axes_styles\": {\n", - " \"line_color\": \"black\",\n", - " \"line_width\": 0.2,\n", - " \"grid_width\": 1,\n", - " \"format_numbers_as_k_m\": True,\n", - " \"exclude_non_numeric_formatting\": True\n", - " },\n", - " \"title\": {\"bold_html\": True, \"include\": True},\n", - " \"default_size\": {\"height\": 1200, \"width\": 1000}\n", - " }\n", - " },\n", - " {\n", - " \"category\": \"histogram_distribution\",\n", - " \"description\": \"Specialized histogram for return distributions with opacity control.\",\n", - " \"styling\": {\n", - " \"template\": \"plotly_white\",\n", - " \"opacity\": 0.75,\n", - " \"axes_styles\": {\n", - " \"grid_width\": 1,\n", - " \"format_numbers_as_k_m\": True,\n", - " \"exclude_non_numeric_formatting\": True\n", - " },\n", - " \"title\": {\"bold_html\": True, \"include\": True},\n", - " \"default_size\": {\"height\": 1200, \"width\": 1000}\n", - " }\n", - " }\n", - "]\n", - "\n", - "# Convert to list of JSON strings\n", - "styling_instructions = [str(chart_dict) for chart_dict in styling_instructions]" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "0f15d1ca", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "[\"{'category': 'line_charts', 'description': 'Used to visualize trends and changes over time, often with multiple series.', 'styling': {'template': 'plotly_white', 'axes_line_width': 0.2, 'grid_width': 1, 'title': {'bold_html': True, 'include': True}, 'colors': 'use multiple colors if more than one line', 'annotations': ['min', 'max'], 'number_format': {'apply_k_m': True, 'thresholds': {'K': 1000, 'M': 100000}, 'percentage_decimals': 2, 'percentage_sign': True}, 'default_size': {'height': 1200, 'width': 1000}}}\",\n", - " \"{'category': 'bar_charts', 'description': 'Useful for comparing discrete categories or groups with bars representing values.', 'styling': {'template': 'plotly_white', 'axes_line_width': 0.2, 'grid_width': 1, 'title': {'bold_html': True, 'include': True}, 'annotations': ['bar values'], 'number_format': {'apply_k_m': True, 'thresholds': {'K': 1000, 'M': 100000}, 'percentage_decimals': 2, 'percentage_sign': True}, 'default_size': {'height': 1200, 'width': 1000}}}\",\n", - " \"{'category': 'histograms', 'description': 'Display the distribution of a data set, useful for returns or frequency distributions.', 'styling': {'template': 'plotly_white', 'bin_size': 50, 'axes_line_width': 0.2, 'grid_width': 1, 'title': {'bold_html': True, 'include': True}, 'annotations': ['x values'], 'number_format': {'apply_k_m': True, 'thresholds': {'K': 1000, 'M': 100000}, 'percentage_decimals': 2, 'percentage_sign': True}, 'default_size': {'height': 1200, 'width': 1000}}}\",\n", - " \"{'category': 'pie_charts', 'description': 'Show composition or parts of a whole with slices representing categories.', 'styling': {'template': 'plotly_white', 'top_categories_to_show': 10, 'bundle_rest_as': 'Others', 'axes_line_width': 0.2, 'grid_width': 1, 'title': {'bold_html': True, 'include': True}, 'annotations': ['x values'], 'number_format': {'apply_k_m': True, 'thresholds': {'K': 1000, 'M': 100000}, 'percentage_decimals': 2, 'percentage_sign': True}, 'default_size': {'height': 1200, 'width': 1000}}}\",\n", - " \"{'category': 'tabular_and_generic_charts', 'description': 'Applies to charts where number formatting needs flexibility, including mixed or raw data.', 'styling': {'template': 'plotly_white', 'axes_line_width': 0.2, 'grid_width': 1, 'title': {'bold_html': True, 'include': True}, 'annotations': ['x values'], 'number_format': {'apply_k_m': True, 'thresholds': {'K': 1000, 'M': 100000}, 'exclude_if_commas_present': True, 'exclude_if_not_numeric': True, 'percentage_decimals': 2, 'percentage_sign': True}, 'default_size': {'height': 1200, 'width': 1000}}}\",\n", - " \"{'category': 'heat_maps', 'description': 'Show data density or intensity using color scales on a matrix or grid.', 'styling': {'template': 'plotly_white', 'axes_styles': {'line_color': 'black', 'line_width': 0.2, 'grid_width': 1, 'format_numbers_as_k_m': True, 'exclude_non_numeric_formatting': True}, 'title': {'bold_html': True, 'include': True}, 'default_size': {'height': 1200, 'width': 1000}}}\",\n", - " \"{'category': 'histogram_distribution', 'description': 'Specialized histogram for return distributions with opacity control.', 'styling': {'template': 'plotly_white', 'opacity': 0.75, 'axes_styles': {'grid_width': 1, 'format_numbers_as_k_m': True, 'exclude_non_numeric_formatting': True}, 'title': {'bold_html': True, 'include': True}, 'default_size': {'height': 1200, 'width': 1000}}}\"]" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "styling_instructions" - ] - }, - { - "cell_type": "code", - "execution_count": 32, - "id": "b39cdaf9", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n" - ] - } - ], - "source": [ - "dspy.configure(lm= dspy.LM('openai/gpt-4o-mini', max_tokens =800, api_key=os.getenv('OPENAI_API_KEY')))\n", - "\n", - "\n", - "print(dspy.settings.lm)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "bcef79e3", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Tables in DuckDB:\n", - "- Month1\n" - ] - }, - { - "data": { - "text/plain": [ - "'| | Month Totals | Xperra | SSC | Michael | Zac | Total | Unnamed: 6 | Unnamed: 7 |\\n|---:|:---------------|---------:|-------:|----------:|------:|--------:|-------------:|-------------:|\\n| 0 | Week1 | 900 | 0 | 0 | 0 | 900 | nan | nan |\\n| 1 | Week2 | 900 | 0 | 0 | 0 | 900 | nan | nan |\\n| 2 | Week3 | 1200 | 0 | 0 | 0 | 1200 | nan | nan |'" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "- Week1\n" - ] - }, - { - "data": { - "text/plain": [ - "'| | Week/Start Min | Unnamed: 1 | Start | 2025-02-24 00:00:00 | End | 2025-03-03 00:00:00 | Unnamed: 6 | Unnamed: 7 | Unnamed: 8 | Unnamed: 9 | Unnamed: 10 |\\n|---:|:-----------------|:-------------|:---------|:----------------------|:------|:----------------------|:-------------|:-------------|:-------------|:-------------|:--------------|\\n| 0 | Start | End | MON | TUE | WED | THU | FRI | SAT | SUN | | |\\n| 1 | 00:00:00 | 00:30:00 | | | | | | | | | |\\n| 2 | 00:30:00 | 01:00:00 | | | | | | | | | |'" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "- Week2\n" - ] - }, - { - "data": { - "text/plain": [ - "'| | Week/Start Min | Unnamed: 1 | Start | 2025-03-03 00:00:00 | End | 2025-03-10 00:00:00 | Unnamed: 6 | Unnamed: 7 | Unnamed: 8 | Unnamed: 9 | Unnamed: 10 |\\n|---:|:-----------------|:-------------|:---------|:----------------------|:------|:----------------------|:-------------|:-------------|:-------------|:-------------|:--------------|\\n| 0 | Start | End | MON | TUE | WED | THU | FRI | SAT | SUN | | |\\n| 1 | 00:00:00 | 00:30:00 | | | | | | | | | |\\n| 2 | 00:30:00 | 01:00:00 | | | | | | | | | |'" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "- Week3\n" - ] - }, - { - "data": { - "text/plain": [ - "'| | Week/Start Min | Unnamed: 1 | Start | 2025-03-10 00:00:00 | End | 2025-03-17 00:00:00 | Unnamed: 6 | Unnamed: 7 | Unnamed: 8 | Unnamed: 9 | Unnamed: 10 |\\n|---:|:-----------------|:-------------|:---------|:----------------------|:------|:----------------------|:-------------|:-------------|:-------------|:-------------|:--------------|\\n| 0 | Start | End | MON | TUE | WED | THU | FRI | SAT | SUN | | |\\n| 1 | 00:00:00 | 00:30:00 | | | | | | | | | |\\n| 2 | 00:30:00 | 01:00:00 | | | | | | | | | |'" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "- Week4\n" - ] - }, - { - "data": { - "text/plain": [ - "'| | Week/Start Min | Unnamed: 1 | Start | 2025-03-17 00:00:00 | End | 2025-03-24 00:00:00 | Unnamed: 6 | Unnamed: 7 | Unnamed: 8 | Unnamed: 9 | Unnamed: 10 |\\n|---:|:-----------------|:-------------|:---------|:----------------------|:------|:----------------------|:-------------|:-------------|:-------------|:-------------|:--------------|\\n| 0 | Start | End | MON | TUE | WED | THU | FRI | SAT | SUN | | |\\n| 1 | 00:00:00 | 00:30:00 | | | | | | | | | |\\n| 2 | 00:30:00 | 01:00:00 | | | | | | | | | |'" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n" - ] - }, - { - "data": { - "text/plain": [ - "('Week4',)" - ] - }, - "execution_count": 19, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "\n", - "excel_path = 'My Timesheet 24th Feb.xlsx'\n", - "\n", - "sheet_names = pd.ExcelFile(excel_path).sheet_names\n", - "\n", - "conn = duckdb.connect()\n", - "\n", - "\n", - "# Create tables for each sheet\n", - "for sheet in sheet_names:\n", - " df = pd.read_excel(excel_path, sheet_name=sheet)\n", - " # Register each DataFrame as a table named after the sheet\n", - " conn.register(sheet, df)\n", - "\n", - "\n", - "# Show all tables in DuckDB\n", - "tables = conn.execute(\"SHOW TABLES\").fetchall()\n", - "print(\"Tables in DuckDB:\")\n", - "for table in tables:\n", - " # Get the first few rows of each table to show structure\n", - " try:\n", - " head_data = conn.execute(f\"SELECT * FROM {table[0]} LIMIT 3\").df().to_markdown()\n", - " display(head_data)\n", - "\n", - " except Exception as e:\n", - " print(f\" Error fetching head of {table[0]}: {e}\")\n", - " print() # Add blank line for readability\n", - "\n", - "\n", - "\n", - "\n", - "# Read Excel file using DuckDB\n", - "# excel = duckdb.sql(\"SELECT * FROM read_excel('My Timesheet 24th Feb.xlsx')\")\n", - "\n", - "\n", - "\n", - "# help(excel)\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "# Preprocessing steps\n", - "# 1. Drop empty rows and columns\n", - "# excel_df.dropna(how='all', inplace=True) # Remove empty rows\n", - "# excel_df.dropna(how='all', axis=1, inplace=True) # Remove empty columns\n", - "\n", - "# # 2. Clean column names\n", - "# excel_df.columns = excel_df.columns.str.strip() # Remove extra spaces\n", - "\n", - "# # 3. Convert Excel data to CSV with UTF-8-sig encoding\n", - "# csv_buffer = io.StringIO()\n", - "# excel_df.to_csv(csv_buffer, index=False, encoding='utf-8-sig')\n", - "# csv_buffer.seek(0)\n", - "\n", - "# # Read the processed CSV back into a dataframe\n", - "# new_df = pd.read_csv(csv_buffer)\n", - "\n", - "\n", - "# excel\n", - "\n", - "table" - ] - }, - { - "cell_type": "code", - "execution_count": 61, - "id": "abf0addd", - "metadata": {}, - "outputs": [], - "source": [ - "class data_context_gen(dspy.Signature):\n", - " \"\"\"\n", - " Generate a compact JSON data context for DuckDB tables ingested from Excel or CSV files.\n", - " The JSON must include:\n", - " - Exact DuckDB table names\n", - " - Source sheet or file name for each table\n", - " - Table role (fact/dimension)\n", - " - Primary key (pk)\n", - " - Columns with type and role (pk, fk, attr, cat, measure, temporal)\n", - " - Relationships between tables (foreign keys), with cardinality types (1:1, 1:M, M:1, M:M)\n", - " - Business purpose of each table\n", - " - Metrics expressed as formulas\n", - " - Use cases for the dataset\n", - "\n", - " Example JSON format:\n", - " {\n", - " \"tables\": {\n", - " \"customer_master\": {\n", - " \"source\": \"Customer_Master sheet\",\n", - " \"role\": \"dimension\",\n", - " \"pk\": \"customer_id\",\n", - " \"columns\": {\n", - " \"customer_id\": {\"type\": \"string\", \"role\": \"pk\"},\n", - " \"name\": {\"type\": \"string\", \"role\": \"attr\"},\n", - " \"region\": {\"type\": \"string\", \"role\": \"cat\"},\n", - " \"signup_date\": {\"type\": \"date\", \"role\": \"temporal\"}\n", - " },\n", - " \"purpose\": \"Customer attributes for segmentation\"\n", - " },\n", - " \"sales_data\": {\n", - " \"source\": \"Sales_Data sheet\",\n", - " \"role\": \"fact\",\n", - " \"pk\": \"order_id\",\n", - " \"columns\": {\n", - " \"order_id\": {\"type\": \"string\", \"role\": \"pk\"},\n", - " \"customer_id\": {\"type\": \"string\", \"role\": \"fk\"},\n", - " \"product_id\": {\"type\": \"string\", \"role\": \"fk\"},\n", - " \"order_date\": {\"type\": \"date\", \"role\": \"temporal\"},\n", - " \"quantity\": {\"type\": \"int\", \"role\": \"measure\"},\n", - " \"unit_price\": {\"type\": \"decimal\", \"role\": \"measure\"}\n", - " },\n", - " \"purpose\": \"Transaction records for revenue analysis\"\n", - " },\n", - " \"product_catalog\": {\n", - " \"source\": \"Product_Catalog sheet\",\n", - " \"role\": \"dimension\",\n", - " \"pk\": \"product_id\",\n", - " \"columns\": {\n", - " \"product_id\": {\"type\": \"string\", \"role\": \"pk\"},\n", - " \"product_name\": {\"type\": \"string\", \"role\": \"attr\"},\n", - " \"category\": {\"type\": \"string\", \"role\": \"cat\"},\n", - " \"subcategory\": {\"type\": \"string\", \"role\": \"cat\"},\n", - " \"brand\": {\"type\": \"string\", \"role\": \"cat\"}\n", - " },\n", - " \"purpose\": \"Product hierarchy for analysis\"\n", - " }\n", - " },\n", - " \"relationships\": [\n", - " {\"from\": \"sales_data.customer_id\", \"to\": \"customer_master.customer_id\", \"type\": \"M:1\"},\n", - " {\"from\": \"sales_data.product_id\", \"to\": \"product_catalog.product_id\", \"type\": \"M:1\"}\n", - " ],\n", - " \"metrics\": [\n", - " \"revenue = quantity * unit_price\",\n", - " \"customer_lifetime_value\"\n", - " ],\n", - " \"use_cases\": [\n", - " \"cohort analysis\",\n", - " \"product performance\",\n", - " \"regional sales\"\n", - " ]\n", - " }\n", - "\n", - " Column roles: pk (primary key), fk (foreign key), attr (attribute), cat (categorical), measure (numerical), temporal (date/time)\n", - " Table roles: fact (transactional), dimension (reference data)\n", - " Relationship types: 1:1, 1:M, M:1, M:M\n", - " \"\"\"\n", - " user_description = dspy.InputField(desc=\"User's description of the data, including relationships\")\n", - " dataset_view = dspy.InputField(desc=\"Dataset name with sample head(5 rows) view\")\n", - " data_context = dspy.OutputField(desc=\"Compact JSON describing DuckDB tables, columns, relationships, metrics and use cases\")" - ] - }, - { - "cell_type": "code", - "execution_count": 71, - "id": "59699a12", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "Prediction(\n", - " data_context='{\\n \"tables\": {\\n \"Month1\": {\\n \"source\": \"Month1 sheet\",\\n \"role\": \"fact\",\\n \"pk\": \"week\",\\n \"columns\": {\\n \"week\": {\"type\": \"string\", \"role\": \"pk\"},\\n \"month_totals\": {\"type\": \"int\", \"role\": \"measure\"},\\n \"xperra\": {\"type\": \"int\", \"role\": \"measure\"},\\n \"ssc\": {\"type\": \"int\", \"role\": \"measure\"},\\n \"michael\": {\"type\": \"int\", \"role\": \"measure\"},\\n \"zac\": {\"type\": \"int\", \"role\": \"measure\"},\\n \"total\": {\"type\": \"int\", \"role\": \"measure\"}\\n },\\n \"purpose\": \"Monthly totals for performance tracking\"\\n },\\n \"Week1\": {\\n \"source\": \"Week1 sheet\",\\n \"role\": \"dimension\",\\n \"pk\": \"week_start\",\\n \"columns\": {\\n \"week_start\": {\"type\": \"datetime\", \"role\": \"pk\"},\\n \"start\": {\"type\": \"time\", \"role\": \"temporal\"},\\n \"end\": {\"type\": \"time\", \"role\": \"temporal\"}\\n },\\n \"purpose\": \"Details of the first week for time analysis\"\\n },\\n \"Week2\": {\\n \"source\": \"Week2 sheet\",\\n \"role\": \"dimension\",\\n \"pk\": \"week_start\",\\n \"columns\": {\\n \"week_start\": {\"type\": \"datetime\", \"role\": \"pk\"},\\n \"start\": {\"type\": \"time\", \"role\": \"temporal\"},\\n \"end\": {\"type\": \"time\", \"role\": \"temporal\"}\\n },\\n \"purpose\": \"Details of the second week for time analysis\"\\n },\\n \"Week3\": {\\n \"source\": \"Week3 sheet\",\\n \"role\": \"dimension\",\\n \"pk\": \"week_start\",\\n \"columns\": {\\n \"week_start\": {\"type\": \"datetime\", \"role\": \"pk\"},\\n \"start\": {\"type\": \"time\", \"role\": \"temporal\"},\\n \"end\": {\"type\": \"time\", \"role\": \"temporal\"}\\n },\\n \"purpose\": \"Details of the third week for time analysis\"\\n },\\n \"Week4\": {\\n \"source\": \"Week4 sheet\",\\n \"role\": \"dimension\",\\n \"pk\": \"week_start\",\\n \"columns\": {\\n \"week_start\": {\"type\": \"datetime\", \"role\": \"pk\"},\\n \"start\": {\"type\": \"time\", \"role\": \"temporal\"},\\n \"end\": {\"type\": \"time\", \"role\": \"temporal\"}\\n },\\n \"purpose\": \"Details of the fourth week for time analysis\"\\n }\\n },\\n \"relationships\": [\\n {\"from\": \"Month1.week\", \"to\": \"Week1.week_start\", \"type\": \"M:1\"},\\n {\"from\": \"Month1.week\", \"to\": \"Week2.week_start\", \"type\": \"M:1\"},\\n {\"from\": \"Month1.week\", \"to\": \"Week3.week_start\", \"type\": \"M:1\"},\\n {\"from\": \"Month1.week\", \"to\": \"Week4.week_start\", \"type\": \"M:1\"}\\n ],\\n \"metrics\": [\\n \"total_monthly_performance = SUM(month_totals)\",\\n \"average_weekly_performance = AVG(month_totals)\"\\n ],\\n \"use_cases\": [\\n \"monthly performance tracking\",\\n \"weekly trend analysis\",\\n \"resource allocation planning\"\\n ]\\n}'\n", - ")" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "user_description = \"These are my worksheets month over month\"\n", - "\n", - "data_context_agent = dspy.Predict(data_context_gen)\n", - "\n", - "tables = conn.execute(\"SHOW TABLES\").fetchall()\n", - "\n", - "dataset_view = \"\"\n", - "\n", - "for table in tables:\n", - " head_data = conn.execute(f\"SELECT * FROM {table[0]} LIMIT 3\").df().to_markdown()\n", - "\n", - " dataset_view+=\"exact_table_name=\"+table[0]+'\\n:'+head_data+'\\n'\n", - "\n", - "\n", - "response = data_context_agent(user_description=user_description, dataset_view=dataset_view)\n", - "\n", - "\n", - "display(response)\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7ec51128", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{\n", - " \"tables\": {\n", - " \"Month1\": {\n", - " \"source\": \"Month1 sheet\",\n", - " \"role\": \"fact\",\n", - " \"pk\": \"week\",\n", - " \"columns\": {\n", - " \"week\": {\"type\": \"string\", \"role\": \"pk\"},\n", - " \"month_totals\": {\"type\": \"int\", \"role\": \"measure\"},\n", - " \"xperra\": {\"type\": \"int\", \"role\": \"measure\"},\n", - " \"ssc\": {\"type\": \"int\", \"role\": \"measure\"},\n", - " \"michael\": {\"type\": \"int\", \"role\": \"measure\"},\n", - " \"zac\": {\"type\": \"int\", \"role\": \"measure\"},\n", - " \"total\": {\"type\": \"int\", \"role\": \"measure\"}\n", - " },\n", - " \"purpose\": \"Monthly totals for performance tracking\"\n", - " },\n", - " \"Week1\": {\n", - " \"source\": \"Week1 sheet\",\n", - " \"role\": \"dimension\",\n", - " \"pk\": \"week_start\",\n", - " \"columns\": {\n", - " \"week_start\": {\"type\": \"datetime\", \"role\": \"pk\"},\n", - " \"start\": {\"type\": \"time\", \"role\": \"temporal\"},\n", - " \"end\": {\"type\": \"time\", \"role\": \"temporal\"}\n", - " },\n", - " \"purpose\": \"Details of the first week for time analysis\"\n", - " },\n", - " \"Week2\": {\n", - " \"source\": \"Week2 sheet\",\n", - " \"role\": \"dimension\",\n", - " \"pk\": \"week_start\",\n", - " \"columns\": {\n", - " \"week_start\": {\"type\": \"datetime\", \"role\": \"pk\"},\n", - " \"start\": {\"type\": \"time\", \"role\": \"temporal\"},\n", - " \"end\": {\"type\": \"time\", \"role\": \"temporal\"}\n", - " },\n", - " \"purpose\": \"Details of the second week for time analysis\"\n", - " },\n", - " \"Week3\": {\n", - " \"source\": \"Week3 sheet\",\n", - " \"role\": \"dimension\",\n", - " \"pk\": \"week_start\",\n", - " \"columns\": {\n", - " \"week_start\": {\"type\": \"datetime\", \"role\": \"pk\"},\n", - " \"start\": {\"type\": \"time\", \"role\": \"temporal\"},\n", - " \"end\": {\"type\": \"time\", \"role\": \"temporal\"}\n", - " },\n", - " \"purpose\": \"Details of the third week for time analysis\"\n", - " },\n", - " \"Week4\": {\n", - " \"source\": \"Week4 sheet\",\n", - " \"role\": \"dimension\",\n", - " \"pk\": \"week_start\",\n", - " \"columns\": {\n", - " \"week_start\": {\"type\": \"datetime\", \"role\": \"pk\"},\n", - " \"start\": {\"type\": \"time\", \"role\": \"temporal\"},\n", - " \"end\": {\"type\": \"time\", \"role\": \"temporal\"}\n", - " },\n", - " \"purpose\": \"Details of the fourth week for time analysis\"\n", - " }\n", - " },\n", - " \"relationships\": [\n", - " {\"from\": \"Month1.week\", \"to\": \"Week1.week_start\", \"type\": \"M:1\"},\n", - " {\"from\": \"Month1.week\", \"to\": \"Week2.week_start\", \"type\": \"M:1\"},\n", - " {\"from\": \"Month1.week\", \"to\": \"Week3.week_start\", \"type\": \"M:1\"},\n", - " {\"from\": \"Month1.week\", \"to\": \"Week4.week_start\", \"type\": \"M:1\"}\n", - " ],\n", - " \"metrics\": [\n", - " \"total_monthly_performance = SUM(month_totals)\",\n", - " \"average_weekly_performance = AVG(month_totals)\"\n", - " ],\n", - " \"use_cases\": [\n", - " \"monthly performance tracking\",\n", - " \"weekly trend analysis\",\n", - " \"resource allocation planning\"\n", - " ]\n", - "}\n" - ] - } - ], - "source": [ - "# print(response.data_context)\n", - "\n", - "l1 = ['muhammadalisaif2@gmail.com',\n", - "faisalnazeer67@gmail.com\n", - "muhammadhunainzubair@gmail.com\n", - "nabeeluddin266@gmail.com\n", - "Mohidyamin@gmail.com\n", - "daemazeemdean@gmail.com\n", - "fariakhaliq123@gmail.com\n", - "ibrahimarain15@gmail.com\n", - "umairahmed805805@gmail.com\n", - "muhibali369@gmail.com\n", - "taharizwan492@gmail.com\n", - "qayyumahmed552@gmail.com\n", - "abdulmusawir3545@gmail.com\n", - "huzaifa.m.awan@gmail.com\n", - "bilalhassan2103@gmail.com\n", - "muhammadabbas485@gmail.com\n", - "eros1030109@gmail.com\n", - "muhammadhureran8@gmail.com\n", - "abdulmoezmughl@gmail.com\n", - "samahabatool7@gmail.com\n", - "3meraldgg@gmail.com\n", - "abdurrehman.azeem81@hotmail.com\n", - "ayesh.bangash810@gmail.com\n", - "'umernaeem.12513@gmail.com']" - ] - }, - { - "cell_type": "code", - "execution_count": 42, - "id": "31ac856e", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Month1\n", - ":| | Month Totals | Xperra | SSC | Michael | Zac | Total | Unnamed: 6 | Unnamed: 7 |\n", - "|---:|:---------------|---------:|-------:|----------:|------:|--------:|-------------:|-------------:|\n", - "| 0 | Week1 | 900 | 0 | 0 | 0 | 900 | nan | nan |\n", - "| 1 | Week2 | 900 | 0 | 0 | 0 | 900 | nan | nan |\n", - "| 2 | Week3 | 1200 | 0 | 0 | 0 | 1200 | nan | nan |\n", - "Week1\n", - ":| | Week/Start Min | Unnamed: 1 | Start | 2025-02-24 00:00:00 | End | 2025-03-03 00:00:00 | Unnamed: 6 | Unnamed: 7 | Unnamed: 8 | Unnamed: 9 | Unnamed: 10 |\n", - "|---:|:-----------------|:-------------|:---------|:----------------------|:------|:----------------------|:-------------|:-------------|:-------------|:-------------|:--------------|\n", - "| 0 | Start | End | MON | TUE | WED | THU | FRI | SAT | SUN | | |\n", - "| 1 | 00:00:00 | 00:30:00 | | | | | | | | | |\n", - "| 2 | 00:30:00 | 01:00:00 | | | | | | | | | |\n", - "Week2\n", - ":| | Week/Start Min | Unnamed: 1 | Start | 2025-03-03 00:00:00 | End | 2025-03-10 00:00:00 | Unnamed: 6 | Unnamed: 7 | Unnamed: 8 | Unnamed: 9 | Unnamed: 10 |\n", - "|---:|:-----------------|:-------------|:---------|:----------------------|:------|:----------------------|:-------------|:-------------|:-------------|:-------------|:--------------|\n", - "| 0 | Start | End | MON | TUE | WED | THU | FRI | SAT | SUN | | |\n", - "| 1 | 00:00:00 | 00:30:00 | | | | | | | | | |\n", - "| 2 | 00:30:00 | 01:00:00 | | | | | | | | | |\n", - "Week3\n", - ":| | Week/Start Min | Unnamed: 1 | Start | 2025-03-10 00:00:00 | End | 2025-03-17 00:00:00 | Unnamed: 6 | Unnamed: 7 | Unnamed: 8 | Unnamed: 9 | Unnamed: 10 |\n", - "|---:|:-----------------|:-------------|:---------|:----------------------|:------|:----------------------|:-------------|:-------------|:-------------|:-------------|:--------------|\n", - "| 0 | Start | End | MON | TUE | WED | THU | FRI | SAT | SUN | | |\n", - "| 1 | 00:00:00 | 00:30:00 | | | | | | | | | |\n", - "| 2 | 00:30:00 | 01:00:00 | | | | | | | | | |\n", - "Week4\n", - ":| | Week/Start Min | Unnamed: 1 | Start | 2025-03-17 00:00:00 | End | 2025-03-24 00:00:00 | Unnamed: 6 | Unnamed: 7 | Unnamed: 8 | Unnamed: 9 | Unnamed: 10 |\n", - "|---:|:-----------------|:-------------|:---------|:----------------------|:------|:----------------------|:-------------|:-------------|:-------------|:-------------|:--------------|\n", - "| 0 | Start | End | MON | TUE | WED | THU | FRI | SAT | SUN | | |\n", - "| 1 | 00:00:00 | 00:30:00 | | | | | | | | | |\n", - "| 2 | 00:30:00 | 01:00:00 | | | | | | | | | |\n" - ] - } - ], - "source": [ - "print(dataset_view)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "92a291cf", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": 67, - "id": "83287b47", - "metadata": {}, - "outputs": [], - "source": [ - "class data_maker(dspy.Signature):\n", - " \"\"\"\n", - " Generate DuckDB SQL queries to fetch data from multiple datasets. Handle joins, aggregations, and filtering across tables.\n", - " Use table names as they appear in dataset_descriptions. Common patterns:\n", - " \n", - " Single table: SELECT * FROM customer_master WHERE region = 'North'\n", - " Join tables: SELECT c.name, SUM(s.quantity * s.unit_price) as revenue \n", - " FROM customer_master c JOIN sales_data s ON c.customer_id = s.customer_id\n", - " Multi-dataset: SELECT e.first_name, a.hours_worked FROM employee_info e \n", - " JOIN attendance_log a ON e.emp_id = a.emp_id WHERE a.date = '2024-01-15'\n", - " Aggregation: SELECT category, COUNT(*) as products FROM product_catalog GROUP BY category\n", - " Time-based: SELECT DATE_TRUNC('month', order_date) as month, SUM(quantity) \n", - " FROM sales_data WHERE order_date >= '2024-01-01' GROUP BY month\n", - " \n", - " Always return: df = conn.execute('SQL_QUERY').df() or more\n", - " \"\"\"\n", - " user_query = dspy.InputField(desc=\"what the user is requesting\")\n", - " dataset_descriptions = dspy.InputField(desc=\"Dict of dataset contexts with table names, columns, and relationships\")\n", - " duckdb_sql = dspy.OutputField(desc=\"df = conn.execute('SQL query to fetch the right data').df()\")" - ] - }, - { - "cell_type": "code", - "execution_count": 68, - "id": "b06c9724", - "metadata": {}, - "outputs": [], - "source": [ - "data_maker_agent = dspy.Predict(data_maker)\n", - "\n", - "user_query = \"show me how much I made from Xperra\"\n", - "\n", - "dataset_descriptions = str(response.data_context)\n", - "\n", - "sql = data_maker_agent(user_query=user_query, dataset_descriptions=dataset_descriptions)\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": 69, - "id": "8c6ef476", - "metadata": {}, - "outputs": [], - "source": [ - "exec(sql.duckdb_sql)" - ] - }, - { - "cell_type": "code", - "execution_count": 70, - "id": "4b1e7b92", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " total_xperra\n", - "0 27240.0\n" - ] - } - ], - "source": [ - "print(df)" - ] - }, - { - "cell_type": "code", - "execution_count": 60, - "id": "9b4b3e34", - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
    \n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
    name
    0Month1
    1Week1
    2Week2
    3Week3
    4Week4
    \n", - "
    " - ], - "text/plain": [ - " name\n", - "0 Month1\n", - "1 Week1\n", - "2 Week2\n", - "3 Week3\n", - "4 Week4" - ] - }, - "execution_count": 60, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "conn.execute(\"SHOW TABLES\").df()\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d532107b", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "base", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.7" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -}