Update crew.py and README.md for plot agent

plot_agent/crewai_plot_agent/README.md

@@ -1,22 +1,22 @@
 # SQL and Plot Workflow API
 
-A powerful data visualization system that uses AI agents to automatically generate SQL queries and create meaningful visualizations from natural language requests. Built with CrewAI framework and deployable on TrueFoundry.
+A powerful data visualization system that uses AI agents to automatically generate SQL queries and create meaningful visualizations from natural language requests. Built with the CrewAI framework, integrated with Traceloop for trace monitoring, and deployable on TrueFoundry.
 
 ## Architecture Overview
 
 This project consists of several key components working together:
 
 1. **Query Agent**: An AI agent powered by CrewAI that:
-   - Uses GPT-4o for natural language understanding
+   - Uses a specified LLM model (see `MODEL_NAME` in environment) for natural language understanding
    - Generates appropriate SQL queries for ClickHouse
    - Handles data extraction and preprocessing
    - Validates query safety and performance
 
 2. **Visualization Agent**: A second AI agent that:
-   - Analyzes the data structure and content
+   - Analyzes data structure and content
    - Determines the most appropriate visualization type
    - Generates plots using matplotlib/seaborn
-   - Handles formatting and styling of visualizations
+   - Handles formatting and styling
 
 3. **FastAPI Backend**: RESTful API that:
    - Coordinates between agents
@@ -35,54 +35,64 @@ This project consists of several key components working together:
 2. Query is processed by the Query Agent to generate SQL
 3. SQL is executed against ClickHouse database
 4. Results are passed to Visualization Agent
-5. Visualization Agent creates appropriate plots
+5. Visualization Agent creates plots
 6. Results are displayed in Streamlit UI
 
 ## Setup and Installation
 
 ### Prerequisites
 
-Ensure you have Python >=3.10 <3.13 installed on your system. This project uses [UV](https://docs.astral.sh/uv/) for dependency management.
+Ensure you have Python >=3.10 and <3.13 installed on your system. This project uses [UV](https://docs.astral.sh/uv/) for dependency management.
 
 1. Install UV:
 ```bash
 pip install uv
 ```
 
-
 ### Environment Configuration
 
 1. Copy the example environment file:
 ```bash
 cp .env.example .env
 ```
 
-2. Edit the `.env` file with your actual credentials:
+2. Edit the `.env` file with your actual credentials and required keys (including Traceloop and Model Name):
 ```env
 # Truefoundry LLMGateway Configuration
 LLM_GATEWAY_BASE_URL=your_llm_gateway_base_url_here
 LLM_GATEWAY_API_KEY=your_llm_gateway_api_key_here
 
+# Required virtual truefoundry LLM Model (e.g. chatwithtraces/gpt5, gpt-4o, gpt-3.5-turbo, etc.)
+MODEL_NAME=your_virtual_model_here
+
+# Traceloop (Tracing/Monitoring) Configuration
+TFY_API_KEY=your_tfy_or_llm_gateway_api_key_here
+TRACELOOP_APP_NAME=sql-agent-crewai
+TRACELOOP_TRACING_PROJECT_FQN=your_traceloop_project_fqn_here
+
 # ClickHouse Database Configuration
 CLICKHOUSE_HOST=your_clickhouse_host_here
 CLICKHOUSE_PORT=443
 CLICKHOUSE_USER=your_clickhouse_user_here
 CLICKHOUSE_PASSWORD=your_clickhouse_password_here
 CLICKHOUSE_DATABASE=default
 ```
+
 3. Install project dependencies:
 ```bash
 uv sync
 ```
+
 ### How to Get Your LLM Gateway Base URL
 
-You can find your LLM_GATEWAY_BASE_URL from the TrueFoundry console:
-	1.	Navigate to the LLM Gateway service on your TrueFoundry dashboard.
-	2.	Click on the deployed endpoint you want to use. Then click on code in the top right.
-	3.	Copy the base part of the URL from the endpoint details. It usually looks like:
+You can find your `LLM_GATEWAY_BASE_URL` from the TrueFoundry console:
+1. Navigate to the LLM Gateway service on your TrueFoundry dashboard.
+2. Click on the deployed endpoint you want to use, then click on "code" in the top right.
+3. Copy the base part of the URL from the endpoint details. It usually looks like:
  ```
 https://<your-platform-url>/api/llm/api/inference/openai
 ```
+
 ### Configuration
 
 1. **Agent Configuration**:
@@ -94,8 +104,7 @@ https://<your-platform-url>/api/llm/api/inference/openai
    - Set task descriptions and expected outputs
 
 3. **Crew Setup**:
-   - Modify `src/crewai_plot_agent/crew.py` to customize agent tools and logic
-   - Adjust task dependencies and workflow
+   - Modify `src/crewai_plot_agent/crew.py` to customize agent tools, traceloop integration, model logic, and workflow
 
 ### Running the Services
 
@@ -104,7 +113,7 @@ https://<your-platform-url>/api/llm/api/inference/openai
 uv run crewai run
 ```
 
-2. Start the FastAPI server:
+2. Start the FastAPI server (includes Traceloop initialization):
 ```bash
 uv run python -m crewai_plot_agent.api
 ```
@@ -116,38 +125,55 @@ uv run python -m streamlit run crewai_plot_agent/app.py
 
 ## Deployment on TrueFoundry
 
-Follow the same deployment steps as mentioned in the original documentation, with these CrewAI-specific additions:
+When deploying, make sure to add and correctly configure both the LLM model name and Traceloop environment variables:
 
-1. Add CrewAI environment variables to your deployment configuration:
 ```python
 env={
-    "OPENAI_API_KEY": "your_openai_api_key",
+    "MODEL_NAME": "your_virtual_model_name here",                  # Specify the LLM model name explicitly
+    "LLM_GATEWAY_BASE_URL": "https://your-llm-url",       # LLM Gateway URL
+    "LLM_GATEWAY_API_KEY": "your_llm_gateway_api_key",
+    "TFY_API_KEY": "your_tfy_or_llm_gateway_api_key",     # Used by both CrewAI and Traceloop
+    "TRACELOOP_APP_NAME": "sql-agent-crewai",
+    "TRACELOOP_TRACING_PROJECT_FQN": "your_traceloop_project_fqn_here",
     "CLICKHOUSE_HOST": "your_clickhouse_host",
     "CLICKHOUSE_PORT": "443",
     "CLICKHOUSE_USER": "your_user",
     "CLICKHOUSE_PASSWORD": "your_password",
     "CLICKHOUSE_DATABASE": "default",
-    "CREWAI_VERBOSE": "true",  # For detailed CrewAI logs
-    "TRACELOOP_BASE_URL": "<your_host_name>/api/otel" # "https://internal.devtest.truefoundry.tech/api/otel"
-    "TRACELOOP_HEADERS"="Authorization=Bearer%20<your_tfy_api_key>"
-
+    "CREWAI_VERBOSE": "true",                              # For detailed CrewAI logs
+    # Optional override (defaults are set in code):
+    "TRACELOOP_BASE_URL": "https://tfy-eo.truefoundry.cloud/api/otel"
 },
 ```
+*Ensure the same TFY_API_KEY or LLM_GATEWAY_API_KEY is used for both LLM access and Traceloop tracing if possible.*
+
+## Traceloop Integration
+
+This project automatically initializes [Traceloop](https://docs.traceloop.com/) in both the API and CrewAI components, enabling robust LLM tracing and visibility. Initialization occurs via environment variables:
+- `TFY_API_KEY` (required for Traceloop SDK)
+- `TRACELOOP_APP_NAME` (e.g., sql-agent-crewai)
+- `TRACELOOP_TRACING_PROJECT_FQN` (your Traceloop project FQN)
+
+Logs will indicate whether Traceloop initialization succeeded at app and crew startup.
+
+## LLM Model Configuration
+
+Set the environment variable `MODEL_NAME` to specify your target virtual LLM Model. This is required for agent instantiation and passed directly to the CrewAI LLM configuration.
 
 ## Support and Resources
 
-For support with CrewAI specific features:
+For support with CrewAI and Traceloop integration:
 - Visit the [CrewAI documentation](https://docs.crewai.com)
 - Join the [CrewAI Discord](https://discord.com/invite/X4JWnZnxPb)
 - Chat with [CrewAI docs](https://chatg.pt/DWjSBZn)
 
-For general project support:
+General project support:
 - Check the [API documentation](/docs)
 - Monitor application health at `/health`
 - View metrics at `/metrics`
 
 ## API Usage
 
-The API endpoints remain the same as in the original documentation. Refer to the FastAPI documentation at `/docs` for detailed API specifications.
+The API endpoints are unchanged from the original documentation. For full details, refer to the FastAPI documentation at `/docs`. Traceloop visibility for requests is automatic if the required environment variables are set.
 
-For CrewAI-specific debugging and monitoring, enable verbose mode in your configuration or use the `CREWAI_VERBOSE` environment variable.
+For extra debugging, enable verbose logging by setting the `CREWAI_VERBOSE` environment variable to `"true"`.

plot_agent/crewai_plot_agent/crewai_plot_agent/crew.py

@@ -1,17 +1,74 @@
 from typing import Any, Dict, List, Optional
 
-from crewai import Agent, Crew, Process, Task
+from crewai import Agent, Crew, LLM, Process, Task
 from crewai.project import CrewBase, agent, crew, task
 from crewai_plot_agent.tools.clickhouse_tool import ClickHouseTool
 from crewai_plot_agent.tools.plot_tool import PlotTools
 from pydantic import BaseModel, Field
+import os
+import litellm
 
 # If you want to run a snippet of code before or after the crew starts,
 # you can use the @before_kickoff and @after_kickoff decorators
 # https://docs.crewai.com/concepts/crews#example-crew-class-with-decorators
+
+
+
 from traceloop.sdk import Traceloop
 
-Traceloop.init(app_name="crewai")
+# Traceloop Configuration
+TRACELOOP_APP_NAME = os.environ.get("TRACELOOP_APP_NAME")
+TRACELOOP_TRACING_PROJECT_FQN = os.environ.get("TRACELOOP_TRACING_PROJECT_FQN")
+
+# Get API Key for both LLM Gateway and Traceloop
+# TFY_API_KEY can be used for both, or use separate keys
+TFY_API_KEY = os.environ.get("TFY_API_KEY") or os.environ.get("LLM_GATEWAY_API_KEY")
+LLM_GATEWAY_BASE_URL = os.environ.get("LLM_GATEWAY_BASE_URL", "https://gateway.truefoundry.ai")
+MODEL_NAME =os.environ.get("MODEL_NAME")
+
+
+# Validate that required environment variables are set
+if not TFY_API_KEY:
+    raise ValueError(
+        "LLM_GATEWAY_API_KEY must be set. "
+        "Please set LLM_GATEWAY_API_KEY environment variable."
+    )
+
+# Configure LiteLLM default headers
+# litellm.drop_params = True  # Drop unsupported params
+# litellm.suppress_debug_info = False
+
+extra_headers = {
+    "Authorization": f"Bearer {TFY_API_KEY}",
+    "TFY-Tracing-Project": TRACELOOP_TRACING_PROJECT_FQN,
+}
+
+# Set default headers for all LiteLLM requests
+if not hasattr(litellm, 'default_headers') or litellm.default_headers is None:
+    litellm.default_headers = {}
+litellm.default_headers.update(extra_headers)
+
+# Initialize Traceloop SDK for tracing
+# Generate a Personal Access Token or Virtual Account from the Access Tab
+# If you are using Virtual Account, make sure to give it access to the tracing project
+print("\n" + "="*80)
+print("INITIALIZING TRACELOOP SDK FOR TRACING")
+print("="*80)
+if TFY_API_KEY:
+    Traceloop.init(
+        app_name=TRACELOOP_APP_NAME,
+        api_endpoint="https://tfy-eo.truefoundry.cloud/api/otel",
+        headers={
+            "Authorization": f"Bearer {TFY_API_KEY}",
+            "TFY-Tracing-Project": TRACELOOP_TRACING_PROJECT_FQN,
+        },
+    )
+    print("✅ Traceloop SDK initialized successfully in crew.py")
+    # print("="*80 + "\n")
+else:
+    print("⚠️  ERROR: Traceloop SDK not initialized - TFY_API_KEY or LLM_GATEWAY_API_KEY not set")
+    print("Please set TFY_API_KEY or LLM_GATEWAY_API_KEY environment variable")
+    print("="*80 + "\n")
 
 
 class SQLQueryResult(BaseModel):
@@ -53,17 +110,40 @@ class CrewaiPlotAgent:
     # 	print("Output", output ,"sdfsdfsdfsd")
     # 	return output
 
+    # Configure LLM for agents with TrueFoundry Gateway settings
+    def _get_llm(self) -> LLM:
+        """Get configured LLM for agents with TrueFoundry Gateway"""
+        # Use the actual model name directly: chatwithtraces/gpt5
+        return LLM(
+            model=MODEL_NAME,  # Use actual model: chatwithtraces/gpt5
+            api_key=TFY_API_KEY,
+            base_url=LLM_GATEWAY_BASE_URL,
+            custom_llm_provider="openai",  # <— LiteLLM uses this
+            extra_headers={
+            "Authorization": f"Bearer {TFY_API_KEY}",
+            "TFY-Tracing-Project": TRACELOOP_TRACING_PROJECT_FQN,
+            },
+            temperature=0.1,  
+            # Using chatwithtraces/gpt5 directly without wrapper
+            # Headers are set via litellm.default_headers
+        )
+
     # If you would like to add tools to your agents, you can learn more about it here:
     # https://docs.crewai.com/concepts/agents#agent-tools
+
     @agent
     def sql_writer(self) -> Agent:
-        return Agent(config=self.agents_config["sql_writer"], verbose=True, tools=[ClickHouseTool()])
+        return Agent(config=self.agents_config["sql_writer"], 
+        verbose=True, 
+        llm=self._get_llm(),
+        tools=[ClickHouseTool()])
 
     @agent
     def plot_writer(self) -> Agent:
         return Agent(
             config=self.agents_config["plot_writer"],
             verbose=True,
+            llm=self._get_llm(),
             tools=[PlotTools()],
             pydantic_output=PlotResult,
         )
@@ -101,3 +181,5 @@ def crew(self) -> Crew:
             # output_pydantic=True,
             # process=Process.hierarchical, # In case you wanna use that instead https://docs.crewai.com/how-to/Hierarchical/
         )
+
+