# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

%pip install --upgrade --quiet google-cloud-aiplatform[agent_engines,adk] a2a-sdk a2a-sdk[http-server] litellm pydantic pydantic-ai fastmcp numpy python-dotenv nest-asyncio arize-phoenix openinference-instrumentation-google-adk arize-phoenix-otel openinference-instrumentation-pydantic-ai opentelemetry-sdk opentelemetry-exporter-otlp opentelemetry-api

import sys

if "google.colab" in sys.modules:
    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

import os

import nest_asyncio
import vertexai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
LOCATION = "us-central1" # @param {type: "string", placeholder: "[your-location]", isTemplate: true}
# fmt: on

# Create the bucket
BUCKET_NAME = f"{PROJECT_ID}-agent"
BUCKET_URI = f"gs://{BUCKET_NAME}"

# Set environment variables for ADK
os.environ["GOOGLE_CLOUD_PROJECT"] = PROJECT_ID
os.environ["GOOGLE_CLOUD_LOCATION"] = LOCATION
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "TRUE"

# For notebook async support
nest_asyncio.apply()

# Initiate the client
client = vertexai.Client(project=PROJECT_ID, location=LOCATION)

import os
from pathlib import Path
import random
import uvicorn
import threading
import time
import asyncio
import httpx
from datetime import datetime, timedelta
from textwrap import dedent
import numpy as np
import warnings

warnings.filterwarnings("ignore")

# Pydantic agent
from mcp.server.fastmcp import FastMCP
from pydantic_ai import Agent
from pydantic_ai.models.google import GoogleModel
from pydantic_ai.providers.google import GoogleProvider
from pydantic_ai.mcp import MCPServerStdio
from a2a.types import AgentSkill
from vertexai.preview.reasoning_engines.templates.a2a import create_agent_card
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import TaskUpdater
from a2a.types import TaskState, TextPart, UnsupportedOperationError
from a2a.utils import new_agent_text_message
from a2a.utils.errors import ServerError

# ADK agent
from google.adk.models.lite_llm import litellm
from google.adk.models.lite_llm import LiteLlm
from google.adk.agents import LlmAgent, SequentialAgent
from google.adk import Runner
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
from google.adk.sessions import InMemorySessionService
from google.genai import types
from google.adk.a2a.executor.a2a_agent_executor import (
    A2aAgentExecutor,
    A2aAgentExecutorConfig,
)
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
from a2a.server.apps import A2AStarletteApplication
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from a2a.types import TransportProtocol
from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH
from google.adk.tools.agent_tool import AgentTool

# Agent deployment
import vertexai
from vertexai import agent_engines
from vertexai.preview.reasoning_engines import A2aAgent
from google.auth import default
from google.auth.credentials import Credentials
from google.auth.transport.requests import Request as AuthRequest
from a2a.client.client import ClientConfig as A2AClientConfig
from a2a.client.client_factory import ClientFactory as A2AClientFactory
from a2a.types import TransportProtocol as A2ATransport

# Observability
from phoenix.otel import register
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from openinference.instrumentation.pydantic_ai import OpenInferenceSpanProcessor
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from openinference.instrumentation.google_adk import GoogleADKInstrumentor
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter

# Phoenix configuration
os.environ["PHOENIX_API_KEY"] = ""  # <---- UPDATE with your PHOENIX API Key
os.environ["PHOENIX_BASE_URL"] = (
    "https://app.phoenix.arize.com/s/ryoung-meta"  # <---- UPDATE with your Phoenix hostname
)
os.environ["PHOENIX_COLLECTOR_ENDPOINT"] = (
    "https://app.phoenix.arize.com/s/ryoung-meta/v1/trace"  # <---- UPDATE with your (append /v1/trace to hostname)
)
os.environ["PHOENIX_PROJECT_NAME"] = "trading-agent"

# Configure the Phoenix tracer
tracer_provider = register(
    project_name=os.environ["PHOENIX_PROJECT_NAME"],  # Default is 'default'
    auto_instrument=True,  # Auto-instrument your app based on installed OI dependencies
)

class MarketDataGenerator:
    """Generate realistic synthetic market data."""

    def __init__(self, seed: int = 42):
        random.seed(seed)
        np.random.seed(seed)

        # Base prices for common symbols
        self.base_prices = {
            "NVDA": 850.0,
            "AAPL": 185.0,
            "GOOGL": 155.0,
            "MSFT": 420.0,
            "TSLA": 245.0,
        }

    def generate_price_series(self, symbol: str, days: int = 30) -> list[dict]:
        """Generate realistic OHLCV price series."""
        base_price = self.base_prices.get(symbol, 100.0)

        prices = [base_price]
        for _ in range(days - 1):
            drift = random.uniform(-0.005, 0.01)
            shock = random.gauss(0, 0.02)
            new_price = prices[-1] * (1 + drift + shock)
            prices.append(max(new_price, 1.0))

        # Generate OHLCV data
        ohlcv_data = []
        start_date = datetime.now() - timedelta(days=days)

        for i, close in enumerate(prices):
            date = start_date + timedelta(days=i)
            intraday_range = close * random.uniform(0.01, 0.03)
            open_price = close + random.uniform(-intraday_range / 2, intraday_range / 2)
            high = max(open_price, close) + random.uniform(0, intraday_range)
            low = min(open_price, close) - random.uniform(0, intraday_range)
            volume = int(random.uniform(50_000_000, 150_000_000))

            ohlcv_data.append(
                {
                    "date": date.strftime("%Y-%m-%d"),
                    "open": round(open_price, 2),
                    "high": round(high, 2),
                    "low": round(low, 2),
                    "close": round(close, 2),
                    "volume": volume,
                }
            )

        return ohlcv_data

    def _calculate_rsi(self, prices: list[float], period: int = 14) -> float:
        """Calculate RSI indicator."""
        if len(prices) < period + 1:
            return 50.0

        deltas = np.diff(prices[-period - 1 :])
        gains = deltas.copy()
        losses = deltas.copy()
        gains[gains < 0] = 0
        losses[losses > 0] = 0
        losses = abs(losses)

        avg_gain = np.mean(gains) if len(gains) > 0 else 0
        avg_loss = np.mean(losses) if len(losses) > 0 else 0.01

        rs = avg_gain / avg_loss if avg_loss != 0 else 100
        rsi = 100 - (100 / (1 + rs))

        return rsi

    def _calculate_macd(self, prices: list[float]) -> tuple:
        """Calculate MACD and signal line."""
        if len(prices) < 26:
            return (0.0, 0.0)

        # Simplified MACD calculation
        fast_ema = np.mean(prices[-12:])
        slow_ema = np.mean(prices[-26:])
        macd = fast_ema - slow_ema
        signal = macd * 0.9

        return (macd, signal)

market_generator = MarketDataGenerator()

# Constants for risk analysis
RSI_OVERBOUGHT_THRESHOLD = 70
RISK_HIGH_THRESHOLD = 60

# Initialize MCP server for Bear Agent tools
bear_mcp = FastMCP("bear-agent-tools")


@bear_mcp.tool()
async def risk_scanner(symbol: str) -> str:
    """Scan for potential downside risks and warning signals.

    Args:
        symbol: Stock symbol to analyze (e.g., NVDA, AAPL)

    Returns:
        Risk analysis report
    """
    # Generate market data and calculate technical indicators
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    # Calculate overall risk score
    risk_score = np.random.uniform(40, 75)

    # Identify specific risks based on technical indicators
    risks = []

    if rsi > RSI_OVERBOUGHT_THRESHOLD:
        risks.append(
            {
                "risk": "Overbought Conditions",
                "severity": "HIGH",
                "description": f"RSI at {rsi:.1f} indicates potential pullback",
                "impact": "-5% to -10%",
            }
        )

    if len(risks) == 0:
        risks.append(
            {
                "risk": "Valuation Concerns",
                "severity": "MEDIUM",
                "description": "P/E ratio elevated vs historical average",
                "impact": "-10% to -15%",
            }
        )

    # Format comprehensive risk report
    separator = "=" * 40
    risk_level = "HIGH" if risk_score > RISK_HIGH_THRESHOLD else "MEDIUM"

    result = dedent(f"""\
        RISK ANALYSIS FOR {symbol}
        {separator}
        Current Price: ${current_price}
        Risk Score: {risk_score:.1f}/100
        Risk Level: {risk_level}

        Identified Risks:
    """)

    for risk in risks:
        result += f"\n[{risk['severity']}] {risk['risk']}"
        result += f"\n   {risk['description']}"
        result += f"\n   Potential Impact: {risk['impact']}\n"

    return result


@bear_mcp.tool()
async def divergence_detector(symbol: str) -> str:
    """Detect bearish divergences and technical weakness.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Divergence analysis report
    """
    prices = market_generator.generate_price_series(symbol, days=30)
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    divergence_score = np.random.uniform(30, 70)
    separator = "=" * 40

    result = dedent(f"""\
        DIVERGENCE ANALYSIS FOR {symbol}
        {separator}
        Divergence Score: {divergence_score:.1f}/100
        RSI: {rsi:.1f}

        Detected Divergences:
        • RSI Bearish Divergence
           Price making highs but RSI not confirming
           Confidence: 75%

        • Volume Divergence
           Declining volume on advances
           Confidence: 70%
    """)

    return result


@bear_mcp.tool()
async def exit_signal_monitor(symbol: str) -> str:
    """Monitor for distribution patterns and exit signals.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Exit signal analysis
    """
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    # Stop loss levels
    stop_aggressive = round(current_price * 0.95, 2)
    stop_moderate = round(current_price * 0.93, 2)
    separator = "=" * 40

    result = dedent(f"""\
        EXIT SIGNAL MONITOR FOR {symbol}
        {separator}
        Current Price: ${current_price}

        Exit Signals:
        [MED] Distribution Pattern
           Heavy selling on up days
           Action: Reduce position size

        Stop Loss Recommendations:
           Aggressive: ${stop_aggressive} (-5%)
           Moderate: ${stop_moderate} (-7%)
    """)

    return result

# Define the Bear Agent's personality and role
bear_system_prompt = (
    "You are a cautious risk analyst focused on identifying potential downside catalysts, "
    "warning signals, and protective strategies. You prioritize capital preservation. "
    "Use the available MCP tools to analyze market risks comprehensively."
)

# Configure Gemini model for Vertex AI
provider = GoogleProvider(vertexai=True)
model = GoogleModel("gemini-2.5-flash", provider=provider)

# Create Pydantic AI agent with MCP tools and tracing enabled
bear_agent = Agent(
    model=model,
    system_prompt=bear_system_prompt,
    tools=[risk_scanner, divergence_detector, exit_signal_monitor],
    retries=2,
    instrument=True,  # Enable automatic tracing
)

async def test_bear_agent():
    # Test query for risk analysis
    query = "Analyze the risks for NVDA stock"
    print(f"Query: {query}")
    print("-" * 60)

    # Run agent and get response
    result = await bear_agent.run(query)
    print("Agent Response:\n")
    print(result.output)

    # Give Phoenix a moment to receive the data
    await asyncio.sleep(2)


# Execute the test
await test_bear_agent()

# Initialize MCP server for Bull agent
bull_mcp = FastMCP("bull-agent-tools")


@bull_mcp.tool()
async def find_breakout_patterns(symbol: str) -> str:
    """Identify bullish breakout patterns and technical setups.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Breakout analysis report
    """
    # Generate prices
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    # Calculate the score
    breakout_score = np.random.uniform(55, 85)

    result = f"""
BREAKOUT PATTERN ANALYSIS FOR {symbol}
{"=" * 40}
Current Price: ${current_price}
Breakout Score: {breakout_score:.1f}/100
Momentum: {"STRONG" if breakout_score > 70 else "MODERATE"}

Bullish Patterns:
[HIGH] Resistance Breakout
   Price breaking above key resistance
   Target: ${round(current_price * 1.08, 2)} (+8%)

[MED] Ascending Triangle
   Higher lows with resistance test
   Target: ${round(current_price * 1.10, 2)} (+10%)
"""

    return result


@bull_mcp.tool()
async def momentum_screener(symbol: str) -> str:
    """Screen for stocks with strong upward momentum.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Momentum analysis report
    """
    # Generate prices
    prices = market_generator.generate_price_series(symbol, days=30)
    closes = [p["close"] for p in prices]

    # Calculate kpis
    rsi = market_generator._calculate_rsi(closes)
    momentum_score = np.random.uniform(60, 90)

    result = f"""
MOMENTUM ANALYSIS FOR {symbol}
{"=" * 40}
Momentum Score: {momentum_score:.1f}/100
Rating: {"VERY STRONG" if momentum_score > 80 else "STRONG"}
Trend: BULLISH

Momentum Factors:
• Healthy RSI at {rsi:.1f} - room to run
• MACD bullish crossover confirmed
• Volume surge - institutions accumulating
• Uptrend pattern intact
"""

    return result


@bull_mcp.tool()
async def entry_signal_detector(symbol: str) -> str:
    """Detect optimal entry points for long positions.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Entry signal analysis
    """
    # Generate prices
    prices = market_generator.generate_price_series(symbol, days=30)

    current_price = prices[-1]["close"]
    entry_quality = np.random.uniform(60, 90)

    result = f"""
ENTRY SIGNAL ANALYSIS FOR {symbol}
{"=" * 40}
Current Price: ${current_price}
Entry Quality: {entry_quality:.1f}/100

Entry Signals:
[HIGH] Pullback to Support
   Quality entry at ${round(current_price * 0.98, 2)}
   Stop Loss: ${round(current_price * 0.95, 2)}
   Risk/Reward: 1:3

Position Sizing:
   Suggested: {"75-100%" if entry_quality > 80 else "50-75%"} of planned position
"""

    return result

# Model configuration
llama_model = "vertex_ai/meta/llama-3.3-70b-instruct-maas"

# Configure LiteLLM to route requests to Vertex AI
litellm.vertex_project = os.environ.get("GOOGLE_CLOUD_PROJECT")
litellm.vertex_location = os.environ.get("GOOGLE_CLOUD_REGION")

# Define the Bull Agent's personality and role
bull_system_instruction = (
    "You are an optimistic market analyst focused on identifying growth opportunities, "
    "bullish patterns, and upside catalysts. You emphasize potential gains and momentum. "
    "Use the available tools to analyze market opportunities comprehensively."
)

# Create ADK agent
bull_agent = LlmAgent(
    name="bull_agent",
    model=LiteLlm(llama_model),  # Using Llama 3.3 for opportunity analysis
    description="Optimistic analyst focused on growth opportunities and bullish signals.",
    instruction=bull_system_instruction,
    tools=[find_breakout_patterns, momentum_screener, entry_signal_detector],
)

async def test_bull_agent():
    """Test the Bull Agent locally using ADK Runner."""
    # Create ADK Runner to manage agent execution
    runner = Runner(
        app_name=bull_agent.name,
        agent=bull_agent,
        session_service=InMemorySessionService(),  # Manages conversations
    )

    # Create a session for the conversation
    session = await runner.session_service.create_session(
        app_name=bull_agent.name,
        user_id="test_user",
        session_id="test_session",
    )

    # Test query for opportunity analysis
    query = "What are the growth opportunities for AAPL stock?"
    print(f"Query: {query}")
    print("-" * 60)

    # Format message in ADK/Gemini format
    content = types.Content(role="user", parts=[types.Part(text=query)])

    # Run agent and capture final response
    final_response = None
    async for event in runner.run_async(
        session_id=session.id, user_id="test_user", new_message=content
    ):
        # Look for the final response event
        if event.is_final_response():
            final_response = event
            break

    # Extract and display the response
    if final_response and final_response.content:
        print("Agent Response:\n")
        for part in final_response.content.parts:
            if hasattr(part, "text") and part.text:
                print(part.text)


# Execute the test
await test_bull_agent()

# Create directory structure for MCP tools
mcp_tools_dir = Path("mcp_tools")
mcp_tools_dir.mkdir(exist_ok=True)

%%writefile $mcp_tools_dir/__init__.py
"""MCP Tools package for trading agents."""

from .market_data import MarketDataGenerator

__all__ = ["MarketDataGenerator"]

%%writefile $mcp_tools_dir/market_data.py
"""Market Data Generator - Creates synthetic market data for testing."""

import random
import numpy as np
from datetime import datetime, timedelta
from typing import List, Dict


class MarketDataGenerator:
    """Generate realistic synthetic market data."""

    def __init__(self, seed: int = 42):
        random.seed(seed)
        np.random.seed(seed)

        # Base prices for common symbols
        self.base_prices = {
            "NVDA": 850.0,
            "AAPL": 185.0,
            "GOOGL": 155.0,
            "MSFT": 420.0,
            "TSLA": 245.0,
        }

    def generate_price_series(self, symbol: str, days: int = 30) -> List[Dict]:
        """Generate realistic OHLCV price series."""
        base_price = self.base_prices.get(symbol, 100.0)

        prices = [base_price]
        for _ in range(days - 1):
            drift = random.uniform(-0.005, 0.01)
            shock = random.gauss(0, 0.02)
            new_price = prices[-1] * (1 + drift + shock)
            prices.append(max(new_price, 1.0))

        # Generate OHLCV data
        ohlcv_data = []
        start_date = datetime.now() - timedelta(days=days)

        for i, close in enumerate(prices):
            date = start_date + timedelta(days=i)
            intraday_range = close * random.uniform(0.01, 0.03)
            open_price = close + random.uniform(-intraday_range/2, intraday_range/2)
            high = max(open_price, close) + random.uniform(0, intraday_range)
            low = min(open_price, close) - random.uniform(0, intraday_range)
            volume = int(random.uniform(50_000_000, 150_000_000))

            ohlcv_data.append({
                "date": date.strftime("%Y-%m-%d"),
                "open": round(open_price, 2),
                "high": round(high, 2),
                "low": round(low, 2),
                "close": round(close, 2),
                "volume": volume
            })

        return ohlcv_data

    def _calculate_rsi(self, prices: List[float], period: int = 14) -> float:
        """Calculate RSI indicator."""
        if len(prices) < period + 1:
            return 50.0

        deltas = np.diff(prices[-period-1:])
        gains = deltas.copy()
        losses = deltas.copy()
        gains[gains < 0] = 0
        losses[losses > 0] = 0
        losses = abs(losses)

        avg_gain = np.mean(gains) if len(gains) > 0 else 0
        avg_loss = np.mean(losses) if len(losses) > 0 else 0.01

        rs = avg_gain / avg_loss if avg_loss != 0 else 100
        rsi = 100 - (100 / (1 + rs))

        return rsi

    def _calculate_macd(self, prices: List[float]) -> tuple:
        """Calculate MACD and signal line."""
        if len(prices) < 26:
            return (0.0, 0.0)

        # Simplified MACD calculation
        fast_ema = np.mean(prices[-12:])
        slow_ema = np.mean(prices[-26:])
        macd = fast_ema - slow_ema
        signal = macd * 0.9  # Simplified signal

        return (macd, signal)

%%writefile $mcp_tools_dir/bear_tools.py
"""Bear Agent MCP Tools - Risk analysis tools."""

import numpy as np
from mcp.server.fastmcp import FastMCP
from market_data import MarketDataGenerator

# Initialize MCP server
mcp = FastMCP("bear-agent-tools")

# Create global market data generator
market_generator = MarketDataGenerator()


@mcp.tool()
async def risk_scanner(symbol: str) -> str:
    """Scan for potential downside risks and warning signals."""
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    risk_score = np.random.uniform(40, 75)

    risks = []
    if rsi > 70:
        risks.append({
            "risk": "Overbought Conditions",
            "severity": "HIGH",
            "description": f"RSI at {rsi:.1f} indicates potential pullback",
            "impact": "-5% to -10%",
        })

    if len(risks) == 0:
        risks.append({
            "risk": "Valuation Concerns",
            "severity": "MEDIUM",
            "description": "P/E ratio elevated vs historical average",
            "impact": "-10% to -15%",
        })

    result = f"""
RISK ANALYSIS FOR {symbol}
{'='*40}
Current Price: ${current_price}
Risk Score: {risk_score:.1f}/100
Risk Level: {"HIGH" if risk_score > 60 else "MEDIUM"}

Identified Risks:
"""

    for risk in risks:
        result += f"\\n[{risk['severity']}] {risk['risk']}"
        result += f"\\n   {risk['description']}"
        result += f"\\n   Potential Impact: {risk['impact']}\\n"

    return result


@mcp.tool()
async def divergence_detector(symbol: str) -> str:
    """Detect bearish divergences and technical weakness."""
    prices = market_generator.generate_price_series(symbol, days=30)
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    divergence_score = np.random.uniform(30, 70)

    result = f"""
DIVERGENCE ANALYSIS FOR {symbol}
{'='*40}
Divergence Score: {divergence_score:.1f}/100
RSI: {rsi:.1f}

Detected Divergences:
• RSI Bearish Divergence
   Price making highs but RSI not confirming
   Confidence: 75%

• Volume Divergence
   Declining volume on advances
   Confidence: 70%
"""

    return result


@mcp.tool()
async def exit_signal_monitor(symbol: str) -> str:
    """Monitor for distribution patterns and exit signals."""
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    stop_aggressive = round(current_price * 0.95, 2)
    stop_moderate = round(current_price * 0.93, 2)

    result = f"""
EXIT SIGNAL MONITOR FOR {symbol}
{'='*40}
Current Price: ${current_price}

Exit Signals:
[MED] Distribution Pattern
   Heavy selling on up days
   Action: Reduce position size

Stop Loss Recommendations:
   Aggressive: ${stop_aggressive} (-5%)
   Moderate: ${stop_moderate} (-7%)
"""

    return result

%%writefile $mcp_tools_dir/bear_mcp_server.py
"""Bear Agent MCP Server - Risk-focused market analysis tools."""

# Import the mcp instance with all registered tools
from bear_tools import mcp

if __name__ == "__main__":
    # Run the MCP server with STDIO transport
    mcp.run(transport="stdio")

def create_bear_agent_card():
    """Create A2A Agent Card for Bear Risk Analyst."""
    # Define the agent's capabilities as A2A skills
    skills = [
        AgentSkill(
            id="risk_analysis",
            name="Risk Factor Scanner",
            description="Identifies potential downside catalysts and risk factors",
            tags=["Risk-Analysis", "Market-Analysis"],
            examples=[
                "What are the key risks for NVDA?",
                "Analyze downside catalysts for tech stocks",
            ],
        ),
        AgentSkill(
            id="divergence_detection",
            name="Divergence Detection",
            description="Finds bearish divergences and technical weakness signals",
            tags=["Technical-Analysis", "Divergence"],
            examples=[
                "Find bearish divergences in AAPL",
            ],
        ),
        AgentSkill(
            id="exit_signals",
            name="Exit Signal Monitoring",
            description="Tracks distribution patterns and exit signals",
            tags=["Exit-Strategy", "Risk-Management"],
            examples=[
                "Monitor exit signals for NVDA",
            ],
        ),
    ]

    # Create A2A agent card for capability advertisement
    return create_agent_card(
        agent_name="Bear Risk Analyst (Pydantic AI + MCP)",
        description=(
            "A cautious risk analyst powered by Pydantic AI, "
            "focused on identifying downside catalysts and warning signals."
        ),
        skills=skills,
    )


# Generate the agent card
bear_agent_card = create_bear_agent_card()

print("Bear Agent Card:")
print(f"   Name: {bear_agent_card.name}")
print(f"   Skills: {len(bear_agent_card.skills)}")

class BearAgentExecutor(AgentExecutor):
    """Agent executor for A2A integration with Bear Agent."""

    def __init__(self):
        # Agent initialized lazily to avoid pickling issues
        self.agent = None

        # Initiate Phoenix register
        self.register = None

    def _init_agent(self):
        """Initialize Pydantic AI agent with MCP tools on deployment."""
        if self.register is None:
            import os

            from phoenix.otel import register

            phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", "default")

            # Configure the Phoenix tracer
            tracer_provider = register(
                project_name=phoenix_project_name, auto_instrument=True
            )

        if self.agent is None:
            import os

            import vertexai
            from pydantic_ai import Agent
            from pydantic_ai.mcp import MCPServerStdio
            from pydantic_ai.models.google import GoogleModel
            from pydantic_ai.providers.google import GoogleProvider

            # Get configuration from environment
            project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
            location = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1")

            # Initialize Vertex AI
            vertexai.init(project=project_id, location=location)

            # System prompt for Bear Agent
            bear_system_prompt = (
                "You are a cautious risk analyst focused on identifying potential downside catalysts, "
                "warning signals, and protective strategies. You prioritize capital preservation. "
                "Use the available MCP tools to analyze market risks comprehensively."
            )

            # Create provider and model
            provider = GoogleProvider(vertexai=True)
            model = GoogleModel("gemini-2.5-flash", provider=provider)

            # Configure MCP server connection
            mcp_server = MCPServerStdio(
                "python", args=["mcp_tools/bear_mcp_server.py"], timeout=60
            )

            # Create Bear Agent
            self.agent = Agent(
                model=model,
                system_prompt=bear_system_prompt,
                toolsets=[mcp_server],
                retries=3,
            )

    async def cancel(self, context: RequestContext, event_queue: EventQueue):
        # Cancellation not supported
        raise ServerError(error=UnsupportedOperationError())

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Execute Bear Agent analysis."""
        # Initialize agent if needed
        if self.agent is None:
            self._init_agent()

        # Extract user query from A2A context
        query = context.get_user_input()
        updater = TaskUpdater(event_queue, context.task_id, context.context_id)

        # Submit task if not already submitted
        if not hasattr(context, "current_task") or not context.current_task:
            await updater.submit()

        # Mark task as actively working
        await updater.start_work()

        try:
            # Update status to show progress
            await updater.update_status(
                TaskState.working, message=new_agent_text_message("Analyzing risks...")
            )

            # Run agent with user query
            result = await self.agent.run(query)

            # Extract result text from Pydantic AI response
            if hasattr(result, "output"):
                result_text = result.output
            else:
                result_text = str(result)

            # Format response
            response = f"""
BEAR RISK ANALYSIS
{"=" * 50}

{result_text}

Analysis completed
"""
            # Add result as artifact and complete task
            await updater.add_artifact([TextPart(text=response)], name="risk_analysis")
            await updater.complete()

        except Exception as e:
            # Mark task as failed on error
            await updater.update_status(
                TaskState.failed,
                message=new_agent_text_message(f"Analysis failed: {e!s}"),
            )

%%writefile $mcp_tools_dir/bull_tools.py
"""Bull Agent MCP Tools - Opportunity analysis tools."""

import numpy as np
from mcp.server.fastmcp import FastMCP
from market_data import MarketDataGenerator

# Initialize MCP server
mcp = FastMCP("bull-agent-tools")

# Create global market data generator
market_generator = MarketDataGenerator()


@mcp.tool()
async def find_breakout_patterns(symbol: str) -> str:
    """Identify bullish breakout patterns and technical setups."""
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    breakout_score = np.random.uniform(55, 85)

    result = f"""
BREAKOUT PATTERN ANALYSIS FOR {symbol}
{'='*40}
Current Price: ${current_price}
Breakout Score: {breakout_score:.1f}/100
Momentum: {"STRONG" if breakout_score > 70 else "MODERATE"}

Bullish Patterns:
[HIGH] Resistance Breakout
   Price breaking above key resistance
   Target: ${round(current_price * 1.08, 2)} (+8%)

[MED] Ascending Triangle
   Higher lows with resistance test
   Target: ${round(current_price * 1.10, 2)} (+10%)
"""

    return result


@mcp.tool()
async def momentum_screener(symbol: str) -> str:
    """Screen for stocks with strong upward momentum."""
    prices = market_generator.generate_price_series(symbol, days=30)
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    momentum_score = np.random.uniform(60, 90)

    result = f"""
MOMENTUM ANALYSIS FOR {symbol}
{'='*40}
Momentum Score: {momentum_score:.1f}/100
Rating: {"VERY STRONG" if momentum_score > 80 else "STRONG"}
Trend: BULLISH

Momentum Factors:
• Healthy RSI at {rsi:.1f} - room to run
• MACD bullish crossover confirmed
• Volume surge - institutions accumulating
• Uptrend pattern intact
"""

    return result


@mcp.tool()
async def entry_signal_detector(symbol: str) -> str:
    """Detect optimal entry points for long positions."""
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    entry_quality = np.random.uniform(60, 90)

    result = f"""
ENTRY SIGNAL ANALYSIS FOR {symbol}
{'='*40}
Current Price: ${current_price}
Entry Quality: {entry_quality:.1f}/100

Entry Signals:
[HIGH] Pullback to Support
   Quality entry at ${round(current_price * 0.98, 2)}
   Stop Loss: ${round(current_price * 0.95, 2)}
   Risk/Reward: 1:3

Position Sizing:
   Suggested: {"75-100%" if entry_quality > 80 else "50-75%"} of planned position
"""

    return result

%%writefile $mcp_tools_dir/bull_mcp_server.py
"""Bull Agent MCP Server - Opportunity-focused market analysis tools."""

# Import the mcp instance with all registered tools
from bull_tools import mcp

if __name__ == "__main__":
    # Run the MCP server with STDIO transport
    mcp.run(transport="stdio")

def create_bull_agent_card():
    """Create A2A Agent Card for Bull Analyst."""
    skills = [
        AgentSkill(
            id="breakout_detection",
            name="Breakout Pattern Detection",
            description="Identify bullish breakout patterns",
            tags=["technical-analysis", "breakouts"],
            examples=["Find breakout patterns for NVDA"],
        ),
        AgentSkill(
            id="momentum_screening",
            name="Momentum Screening",
            description="Screen for stocks with strong momentum",
            tags=["momentum", "screening"],
            examples=["Find high momentum tech stocks"],
        ),
        AgentSkill(
            id="entry_signals",
            name="Entry Signal Detection",
            description="Detect optimal entry points",
            tags=["entry-points", "timing"],
            examples=["When should I buy AAPL?"],
        ),
    ]

    return create_agent_card(
        agent_name="Bull Market Analyst (ADK + MCP)",
        description=(
            "An optimistic analyst powered by Google ADK, "
            "focused on growth opportunities and bullish patterns."
        ),
        skills=skills,
    )


bull_agent_card = create_bull_agent_card()

print("Bull Agent Card:")
print(f"   Name: {bull_agent_card.name}")
print(f"   Skills: {len(bull_agent_card.skills)}")

class BullAgentExecutor(AgentExecutor):
    """Agent executor for Bull Agent."""

    def __init__(self):
        self.agent = None
        self.runner = None
        self.register = None

    def _init_agent(self):
        """Lazy initialization of the Bull Agent and ADK Runner.
        Creates the agent and runner when first needed.
        This happens on Agent Engine after deployment, not during pickling.
        """
        if self.register is None:
            import os

            from phoenix.otel import register

            phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", "default")

            # Configure the Phoenix tracer
            tracer_provider = register(
                project_name=phoenix_project_name, auto_instrument=True
            )

        if self.agent is None:
            import os

            from google.adk.agents import LlmAgent
            from google.adk.models.lite_llm import LiteLlm, litellm
            from google.adk.tools.mcp_tool import StdioConnectionParams
            from google.adk.tools.mcp_tool.mcp_toolset import (
                MCPToolset,
                StdioServerParameters,
            )

            # Configure LiteLLM to route requests to Vertex AI
            litellm.vertex_project = os.environ.get("GOOGLE_CLOUD_PROJECT")
            litellm.vertex_location = os.environ.get("GOOGLE_CLOUD_REGION")

            # Set project and location
            project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
            location = os.environ.get("GOOGLE_CLOUD_REGION")

            # Initialize Vertex AI
            vertexai.init(project=project_id, location=location)

            # Create Bull Agent with MCP tools
            self.agent = LlmAgent(
                model=LiteLlm("vertex_ai/meta/llama-3.3-70b-instruct-maas"),
                name="bull_market_analyst",
                instruction="""You are an optimistic market analyst focused on identifying growth
                opportunities, bullish catalysts, and upside potential. Use the available MCP
                tools to analyze market opportunities comprehensively.""",
                tools=[
                    MCPToolset(
                        connection_params=StdioConnectionParams(
                            server_params=StdioServerParameters(
                                "python",
                                args=["mcp_tools/bull_mcp_server.py"],
                                timeout=60,
                            ),
                        ),
                    )
                ],
            )

        if self.runner is None:
            from google.adk import Runner
            from google.adk.sessions import InMemorySessionService

            self.runner = Runner(
                app_name=self.agent.name,
                agent=self.agent,
                session_service=InMemorySessionService(),
            )

    async def cancel(self, context: RequestContext, event_queue: EventQueue):
        raise ServerError(error=UnsupportedOperationError())

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Execute Bull Agent analysis."""
        if not context.message:
            return

        user_id = (
            context.message.metadata.get("user_id")
            if context.message and context.message.metadata
            else "a2a_user"
        )

        updater = TaskUpdater(event_queue, context.task_id, context.context_id)

        if not hasattr(context, "current_task") or not context.current_task:
            await updater.submit()

        await updater.start_work()

        query = context.get_user_input()

        try:
            await updater.update_status(
                TaskState.working,
                message=new_agent_text_message("Analyzing opportunities..."),
            )

            # Get or create session
            from google.genai import types

            session = await self.runner.session_service.get_session(
                app_name=self.runner.app_name,
                user_id=user_id,
                session_id=context.context_id,
            ) or await self.runner.session_service.create_session(
                app_name=self.runner.app_name,
                user_id=user_id,
                session_id=context.context_id,
            )

            content = types.Content(role="user", parts=[types.Part(text=query)])

            # Run ADK agent
            final_event = None
            async for event in self.runner.run_async(
                session_id=session.id, user_id=user_id, new_message=content
            ):
                if event.is_final_response():
                    final_event = event

            # Extract response
            if final_event and final_event.content and final_event.content.parts:
                response_text = "".join(
                    part.text
                    for part in final_event.content.parts
                    if hasattr(part, "text") and part.text
                )
                if response_text:
                    await updater.add_artifact(
                        [TextPart(text=response_text)],
                        name="opportunity_analysis",
                    )
                    await updater.complete()
                    return

            await updater.update_status(
                TaskState.failed,
                message=new_agent_text_message("Failed to generate response."),
                final=True,
            )

        except Exception as e:
            await updater.update_status(
                TaskState.failed,
                message=new_agent_text_message(f"Analysis failed: {e!s}"),
                final=True,
            )

# Update Bear Agent card
bear_agent_card.url = "http://localhost:8001"
bear_agent_card.preferred_transport = TransportProtocol.jsonrpc

# Update Bull Agent card
bull_agent_card.url = "http://localhost:8002"
bull_agent_card.preferred_transport = TransportProtocol.jsonrpc

def create_bull_agent_a2a_server(agent, agent_card):
    """Create an A2A server for an ADK agent.

    This wraps an ADK agent with A2A protocol handling, making it
    accessible via HTTP endpoints that follow the A2A specification.

    Args:
        agent: The ADK agent instance (LlmAgent, Agent, etc.)
        agent_card: The A2A AgentCard describing the agent's capabilities

    Returns:
        A2AStarletteApplication instance ready to serve via uvicorn
    """
    # Create ADK Runner for the agent
    # The Runner manages agent execution, sessions, and artifacts
    runner = Runner(
        app_name=agent.name,
        agent=agent,
        session_service=InMemorySessionService(),  # Manages conversation state
    )

    # Configure A2A agent executor
    # This bridges ADK agents with the A2A protocol
    config = A2aAgentExecutorConfig()
    executor = A2aAgentExecutor(runner=runner, config=config)

    # Create A2A request handler
    # Handles incoming A2A protocol requests (message:send, get_task, etc.)
    request_handler = DefaultRequestHandler(
        agent_executor=executor,
        task_store=InMemoryTaskStore(),  # Stores task state
    )

    # Create and return A2A Starlette application
    # This is the ASGI app that uvicorn will serve
    return A2AStarletteApplication(agent_card=agent_card, http_handler=request_handler)


async def run_bull_server(agent, agent_card, port):
    """Run a single agent as an A2A server on the specified port."""
    app = create_bull_agent_a2a_server(agent, agent_card)

    # Configure uvicorn server
    config = uvicorn.Config(
        app.build(),  # Build the ASGI application
        host="127.0.0.1",  # localhost
        port=port,
        log_level="warning",  # Quiet output
        loop="none",  # Use the current event loop
    )

    server = uvicorn.Server(config)
    await server.serve()

def create_bear_a2a_server(agent_card):
    """Create A2A server for Pydantic AI Bear Agent.

    Since Bear Agent uses Pydantic AI (not ADK), we create the A2A server
    directly using the BearAgentExecutor we defined earlier.
    """
    request_handler = DefaultRequestHandler(
        agent_executor=BearAgentExecutor(),
        task_store=InMemoryTaskStore(),
    )

    return A2AStarletteApplication(agent_card=agent_card, http_handler=request_handler)


async def run_bear_server(agent_card, port):
    """Run Bear Agent A2A server (Pydantic AI)."""
    app = create_bear_a2a_server(agent_card)

    config = uvicorn.Config(
        app.build(),
        host="127.0.0.1",
        port=port,
        log_level="warning",
        loop="none",
    )

    server = uvicorn.Server(config)
    await server.serve()

async def start_a2a_servers():
    """Start both Bear and Bull agents as A2A servers."""
    # Create tasks for both servers
    # Bear Agent uses Pydantic AI, so it needs custom A2A server
    # Bull Agent uses ADK, so it uses the standard ADK A2A pattern
    tasks = [
        asyncio.create_task(run_bear_server(bear_agent_card, 8001)),
        asyncio.create_task(run_bull_server(bull_agent, bull_agent_card, 8002)),
    ]

    # Give servers time to start
    await asyncio.sleep(2)

    print("   ✓ Bear Agent A2A server: http://127.0.0.1:8001 (Pydantic AI)")
    print("   ✓ Bull Agent A2A server: http://127.0.0.1:8002 (ADK)")

    # Keep servers running
    try:
        await asyncio.gather(*tasks)
    except KeyboardInterrupt:
        print("Shutting down A2A servers...")


def run_servers_in_background():
    """Run A2A servers in a background thread."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(start_a2a_servers())

# Start the A2A servers in background thread
server_thread = threading.Thread(target=run_servers_in_background, daemon=True)
server_thread.start()

# Wait for servers to be ready
time.sleep(3)

# Create remote proxy for Bear Agent
# RemoteA2aAgent discovers capabilities via the agent card endpoint
remote_bear = RemoteA2aAgent(
    name="bear_risk_analyst",
    description="Analyzes risks and warning signals",
    agent_card=f"http://localhost:8001{AGENT_CARD_WELL_KNOWN_PATH}",
)

# Create remote proxy for Bull Agent
remote_bull = RemoteA2aAgent(
    name="bull_market_analyst",
    description="Identifies growth opportunities and bullish patterns",
    agent_card=f"http://localhost:8002{AGENT_CARD_WELL_KNOWN_PATH}",
)

# Create orchestrator that coordinates both agents
trading_orchestrator = LlmAgent(
    name="trading_strategy_orchestrator",
    model="gemini-2.5-flash",
    tools=[
        AgentTool(
            agent=remote_bear,  # Wrap remote agents as tools
        ),
        AgentTool(
            agent=remote_bull,
        ),
    ],
)

# Create Runner for the orchestrator
orchestrator_runner = Runner(
    app_name=trading_orchestrator.name,
    agent=trading_orchestrator,
    session_service=InMemorySessionService(),
)

# Create session
session = await orchestrator_runner.session_service.create_session(
    app_name=trading_orchestrator.name,
    user_id="test_user",
    session_id="orchestrator_test_session",
)

# Test query
test_query = "Should I buy NVDA stock? Analyze opportunity only."

print(f"\nQuery: {test_query}")

# Run orchestrator
content = types.Content(role="user", parts=[types.Part(text=test_query)])

final_result = None
async for event in orchestrator_runner.run_async(
    session_id=session.id, user_id="test_user", new_message=content
):
    if event.is_final_response():
        if event.content and event.content.parts:
            final_result = "".join(
                part.text
                for part in event.content.parts
                if hasattr(part, "text") and part.text
            )
        break

print(f"\nFinal Result:\n{final_result}")

# Configure transport for production deployment
bear_agent_card.preferred_transport = TransportProtocol.http_json

# Wrap agent card and executor in A2A agent
bear_a2a_agent = A2aAgent(
    agent_card=bear_agent_card, agent_executor_builder=BearAgentExecutor
)

# Deploy to Vertex AI Agent Engine
deployed_bear = client.agent_engines.create(
    agent=bear_a2a_agent,
    config={
        "display_name": "Bear Risk Analyst",
        "description": bear_agent_card.description,
        "requirements": [
            "a2a-sdk",
            "google-cloud-aiplatform[agent_engines,adk]",
            "fastmcp",  # Required for MCP tools
            "pydantic",
            "pydantic-ai",  # Required for Bear Agent
            "numpy",
            "arize-phoenix-otel",
            "openinference-instrumentation-pydantic-ai",
            "opentelemetry-sdk",
            "opentelemetry-exporter-otlp",
            "opentelemetry-api",
        ],
        "extra_packages": ["mcp_tools"],  # Include our MCP tools package
        "env_vars": {
            "PHOENIX_API_KEY": os.environ.get("PHOENIX_API_KEY"),
            "PHOENIX_COLLECTOR_ENDPOINT": os.environ.get("PHOENIX_COLLECTOR_ENDPOINT"),
        },
        "staging_bucket": BUCKET_URI,
    },
)

# Configure transport for production deployment
bull_agent_card.preferred_transport = TransportProtocol.http_json

# Create A2A agent
bull_a2a_agent = A2aAgent(
    agent_card=bull_agent_card, agent_executor_builder=BullAgentExecutor
)

# Deploy to Vertex AI Agent Engine
deployed_bull = client.agent_engines.create(
    agent=bull_a2a_agent,
    config={
        "display_name": "Bull Market Analyst",
        "description": bull_agent_card.description,
        "requirements": [
            "a2a-sdk",
            "google-cloud-aiplatform[agent_engines,adk]",
            "fastmcp",  # Required for MCP tools
            "numpy",
            "litellm",
            "arize-phoenix-otel",
            "openinference-instrumentation-google-adk",
        ],
        "env_vars": {
            "PHOENIX_API_KEY": os.environ.get("PHOENIX_API_KEY"),
            "PHOENIX_COLLECTOR_ENDPOINT": os.environ.get("PHOENIX_COLLECTOR_ENDPOINT"),
        },
        "extra_packages": ["mcp_tools"],
        "staging_bucket": BUCKET_URI,
    },
)

# Create GoogleAuth class for httpx authentication
class GoogleAuth(httpx.Auth):
    """Custom httpx Auth class for Google Cloud authentication."""

    def __init__(self) -> None:
        # Get default credentials for the current environment
        self.credentials: Credentials
        self.project: str | None
        self.credentials, self.project = default(
            scopes=["https://www.googleapis.com/auth/cloud-platform"]
        )
        self.auth_request = AuthRequest()

    def auth_flow(self, request: httpx.Request):
        """Add Authorization header to request."""
        # Refresh credentials if expired
        if not self.credentials.valid:
            self.credentials.refresh(self.auth_request)

        # Add Authorization header
        request.headers["Authorization"] = f"Bearer {self.credentials.token}"
        yield request


# Create authenticated httpx client
authenticated_client = httpx.AsyncClient(
    timeout=120,
    auth=GoogleAuth(),  # This adds authentication to ALL requests!
)

# Create client factory for A2A communication
client_config = A2AClientConfig(
    httpx_client=authenticated_client,
    streaming=False,
    polling=False,
    supported_transports=[
        A2ATransport.http_json,
    ],
)

a2a_client_factory = A2AClientFactory(config=client_config)

# Construct Vertex AI Agent Engine API endpoint
api_endpoint = f"https://{LOCATION}-aiplatform.googleapis.com"

# Get resource names from deployed agents
bear_agent_resource_name = deployed_bear.api_resource.name
bull_agent_resource_name = deployed_bull.api_resource.name

# Build A2A endpoint URLs
bear_endpoint = f"{api_endpoint}/v1beta1/{bear_agent_resource_name}/a2a"
bull_endpoint = f"{api_endpoint}/v1beta1/{bull_agent_resource_name}/a2a"

# Create remote agent proxies pointing to deployed endpoints
remote_bear = RemoteA2aAgent(
    name="bear_risk_analyst",
    description="Analyzes risks and warning signals",
    agent_card=f"{bear_endpoint}/v1/card",
    httpx_client=authenticated_client,
    a2a_client_factory=a2a_client_factory,
)

remote_bull = RemoteA2aAgent(
    name="bull_market_analyst",
    description="Identifies growth opportunities and bullish patterns",
    agent_card=f"{bull_endpoint}/v1/card",
    httpx_client=authenticated_client,
    a2a_client_factory=a2a_client_factory,
)

trading_orchestrator = LlmAgent(
    name="trading_strategy_orchestrator",
    model="gemini-2.5-flash",
    tools=[
        AgentTool(
            agent=remote_bear,
        ),
        AgentTool(
            agent=remote_bull,
        ),
    ],
)

# Create Runner for the orchestrator
orchestrator_runner = Runner(
    app_name=trading_orchestrator.name,
    agent=trading_orchestrator,
    session_service=InMemorySessionService(),
)

# Create session
session = await orchestrator_runner.session_service.create_session(
    app_name=trading_orchestrator.name,
    user_id="test_user",
    session_id="orchestrator_test_session",
)

# Test query
test_query = "Analyze the risks for NVDA stock"

print(f"\n📊 Query: {test_query}")

# Run orchestrator
content = types.Content(role="user", parts=[types.Part(text=test_query)])

final_result = None
async for event in orchestrator_runner.run_async(
    session_id=session.id, user_id="test_user", new_message=content
):
    if event.is_final_response():
        if event.content and event.content.parts:
            final_result = "".join(
                part.text
                for part in event.content.parts
                if hasattr(part, "text") and part.text
            )
        break

print(f"\n🤖 Final Result:\n{final_result}")

import os

from phoenix.client import Client
from phoenix.client.types.spans import SpanQuery

# Phoenix configuration
api_key = os.environ["PHOENIX_API_KEY"]
base_url = os.environ["PHOENIX_BASE_URL"]

# Initialize client ## ##Update with base URL and Phoenix API key
client = Client(base_url=base_url, api_key=api_key)

# Get LLM spans for our evaluator
query = (
    SpanQuery()
    .where("span_kind == 'LLM'")
    .select("input.value", "output.value", "llm.tools")
)
orchestrator_df = client.spans.get_spans_dataframe(
    query=query,
    project_identifier=os.environ["PHOENIX_PROJECT_NAME"],
    limit=50,
    timeout=120,
)

TOOL_CALL_PROMPT_TEMPLATE = """
You are an evaluation assistant evaluating user queries and an AI agent's chosen tool calls to
determine whether the tool called would correctly address the user query. The tool
calls have been generated by a AI agent, and chosen from the list of
tools provided below. It is your job to decide whether that agent chose
the right tool to call for the given user query.

    [BEGIN DATA]
    ************
    [User Query]: {input.value}
    ************
    [Tool Called]: {output.value}
    ************
    [Tool Definitions]: {llm.tools}
    ************
    [END DATA]


Your response must be single word, either "correct" or "incorrect",
and should not contain any text or characters aside from that word.
"incorrect" means that the chosen tool would not answer the question,
the tool includes information that is not presented in the question,
or that the tool signature includes parameter values that don't match
the formats specified in the tool signatures below.
"correct" means the correct tool call was chosen, the correct parameters
were extracted from the question, the tool call generated is runnable and correct,
and that no outside information not present in the question was used
in the generated question.

Then write out in a step by step manner an EXPLANATION to show how you determined if the tool selection was correct or incorrect.

EXPLANATION
"""

# Set up and run evaluator
from phoenix.evals import (
    LiteLLMModel,
    llm_classify,
)

os.environ["VERTEXAI_PROJECT"] = os.environ["GOOGLE_CLOUD_PROJECT"]
os.environ["VERTEXAI_LOCATION"] = os.environ["GOOGLE_CLOUD_LOCATION"]
model = LiteLLMModel(model="vertex_ai/gemini-2.5-pro")

# Loop through the dataframe and evaluate each row
evaluations_df = llm_classify(
    dataframe=orchestrator_df,
    template=TOOL_CALL_PROMPT_TEMPLATE,
    model=model,
    rails=["correct", "incorrect"],
    provide_explanation=True,
    concurrency=5,
)

# prep for upload
eval_df = evaluations_df.copy()
eval_df["score"] = eval_df["label"].apply(
    lambda x: 1 if x == "correct" else 0
)  # Create score column
eval_df = eval_df.reset_index()  # Reset the index to make context.span_id a column
eval_df = eval_df.rename(
    columns={"context.span_id": "span_id"}
)  # Rename context.span_id to span_id
eval_df["explanation"] = eval_df["explanation"].str.replace(
    r"^(correct|incorrect)\s*\n*EXPLANATION\s*\n*", "", regex=True
)  # Optional: Clean up explanation text
eval_df = eval_df[
    ["span_id", "label", "score", "explanation"]
]  # Select only the columns you need

# Upload the eval results to Phoenix as annotations.

client.spans.log_span_annotations_dataframe(
    dataframe=eval_df, annotation_name="tool_call_correctness", annotator_kind="LLM"
)

# Get TOOL spans for our evaluator
query = (
    SpanQuery()
    .where("span_kind == 'TOOL'")
    .select("input.value", "output.value", "tool")
)
tools_df = client.spans.get_spans_dataframe(
    query=query,
    project_identifier=os.environ["PHOENIX_PROJECT_NAME"],
    limit=100,
    timeout=120,
)


TOOL_EXECUTION_PROMPT_TEMPLATE = """

You are comparing a function call response to its input and function definition and trying to determine if the generated call response has provided a correct and intended response based on the input. Here is the data:
    [BEGIN DATA]
    ************
    [Input]: {input.value}
    ************
    [Function answer]: {output.value}
    ************
    [END DATA]



Your response must be single word, either "correct" or "incorrect"
and should not contain any text or characters aside from that word.
Compare the input parameters in the generated function against the JSON provided below.
The parameters extracted from the input must match the JSON below exactly.
The function answer should be a correct response to the input.


"correct" means the function call parameters match the JSON below and function answer provides only relevant information.
"incorrect" means that the parameters in the function do not match the JSON schema below exactly, or the function answer does not correctly address the input. You should also respond with "incorrect" if the response makes up information that is not in the JSON schema.

Here are details on the function call:
{tool}

Then write out in a step by step manner an EXPLANATION to show how you determined if the tool selection was correct or incorrect.

EXPLANATION
"""

# Loop through the dataframe and evaluate each row
evaluations_df = llm_classify(
    dataframe=tools_df,
    template=TOOL_EXECUTION_PROMPT_TEMPLATE,
    model=model,
    rails=["correct", "incorrect"],
    provide_explanation=True,
    concurrency=2,
)

# prep for upload
eval_df = evaluations_df.copy()
eval_df["score"] = eval_df["label"].apply(
    lambda x: 1 if x == "correct" else 0
)  # Create score column
eval_df = eval_df.reset_index()  # Reset the index to make context.span_id a column
eval_df = eval_df.rename(
    columns={"context.span_id": "span_id"}
)  # Rename context.span_id to span_id
eval_df["explanation"] = eval_df["explanation"].str.replace(
    r"^(correct|incorrect)\s*\n*EXPLANATION\s*\n*", "", regex=True
)  # Optional: Clean up explanation text
eval_df = eval_df[
    ["span_id", "label", "score", "explanation"]
]  # Select only the columns you need

# Upload the eval results to Phoenix as annotations.

client.spans.log_span_annotations_dataframe(
    dataframe=eval_df,
    annotation_name="tool_execution_correctness",
    annotator_kind="LLM",
)

TRAJECTORY_ACCURACY_PROMPT_WITHOUT_REFERENCE = """
You are a helpful AI bot that checks whether an AI agent’s internal trajectory is accurate and effective.

You will be given:
1. The agent’s actual trajectory of tool calls
2. You will be given input data from a user that the agent used to make a decision
3. You will be given a tool call definition, what the agent used to make the tool call

An accurate trajectory:
- Progresses logically from step to step
- Follows the golden trajectory where reasonable
- Shows a clear path toward completing a goal
- Is reasonably efficient (doesn’t take unnecessary detours)

##

Actual Trajectory:
{tool_calls}

User Inputs:
{attributes.input.value}

Tool Definitions:
{attributes.llm.tools}

##

Your response must be a single string, either `correct` or `incorrect`, and must not include any additional text.

- Respond with `correct` if the agent’s trajectory adheres to the rubric and accomplishes the task effectively.
- Respond with `incorrect` if the trajectory is confusing, misaligned with the goal, inefficient, or does not accomplish the task.

Then write out in a step by step manner an EXPLANATION to show how you determined if the tool selection was correct or incorrect.

EXPLANATION
"""


# Get spans for our evaluator
trajectory_df = client.spans.get_spans_dataframe(
    project_identifier=os.environ["PHOENIX_PROJECT_NAME"], timeout=120
)

# Helper functions for data prep

from typing import Any

import pandas as pd

def filter_spans_by_trace_criteria(
    df: pd.DataFrame,
    trace_filters: dict[str, dict[str, Any]],
    span_filters: dict[str, dict[str, Any]],
) -> pd.DataFrame:
    """Filter spans based on trace-level and span-level criteria.

    Args:
        df: DataFrame with trace data
        trace_filters: Dictionary of column names and filtering criteria for traces
                      Format: {"column_name": {"operator": value}}
                      Supported operators: ">=", "<=", "==", "!=", "contains", "notna", "isna"
        span_filters: Dictionary of column names and filtering criteria for spans
                     Format: {"column_name": {"operator": value}}
                     Same supported operators as trace_filters

    Returns:
        DataFrame with filtered spans from traces that match trace_filters
    """
    # Get all unique trace_ids
    all_trace_ids = set(df["context.trace_id"].unique())
    print(f"Total traces: {len(all_trace_ids)}")

    # Create a copy of the dataframe for filtering
    df_copy = df.copy()

    # Find traces matching the trace criteria
    traces_df = df_copy.copy()
    for column, criteria in trace_filters.items():
        if column not in traces_df.columns:
            print(f"Warning: Column '{column}' not found in dataframe")
            continue

        for operator, value in criteria.items():
            if operator == ">=":
                matching_spans = traces_df[traces_df[column] >= value]
            elif operator == "<=":
                matching_spans = traces_df[traces_df[column] <= value]
            elif operator == "==":
                matching_spans = traces_df[traces_df[column] == value]
            elif operator == "!=":
                matching_spans = traces_df[traces_df[column] != value]
            elif operator == "contains":
                matching_spans = traces_df[
                    traces_df[column].str.contains(value, case=False, na=False)
                ]
            elif operator == "isna":
                matching_spans = traces_df[traces_df[column].isna()]
            elif operator == "notna":
                matching_spans = traces_df[traces_df[column].notna()]
            else:
                print(f"Warning: Unsupported operator '{operator}' - skipping")
                continue

            traces_df = matching_spans

    matching_trace_ids = set(traces_df["context.trace_id"].unique())
    print(f"Found {len(matching_trace_ids)} traces matching trace criteria")

    if not matching_trace_ids:
        print("No matching traces found")
        return pd.DataFrame()

    # Filter to keep only rows from matching traces
    result_df = df[df["context.trace_id"].isin(matching_trace_ids)].copy()

    # Apply span filters
    for column, criteria in span_filters.items():
        if column not in result_df.columns:
            print(f"Warning: Column '{column}' not found in dataframe")
            continue

        for operator, value in criteria.items():
            if operator == ">=":
                result_df = result_df[result_df[column] >= value]
            elif operator == "<=":
                result_df = result_df[result_df[column] <= value]
            elif operator == "==":
                result_df = result_df[result_df[column] == value]
            elif operator == "!=":
                result_df = result_df[result_df[column] != value]
            elif operator == "contains":
                result_df = result_df[
                    result_df[column].str.contains(value, case=False, na=False)
                ]
            elif operator == "isna":
                result_df = result_df[result_df[column].isna()]
            elif operator == "notna":
                result_df = result_df[result_df[column].notna()]
            else:
                print(f"Warning: Unsupported operator '{operator}' - skipping")
                continue

    print(f"Final result: {len(result_df)} spans from {len(matching_trace_ids)} traces")
    return result_df


def prepare_trace_data_for_evaluation(
    df,
    group_by_col="context.trace_id",
    extract_cols={"tool_calls": "tool_calls"},
    additional_data=None,
    filter_empty=True,
):
    """Prepare trace data for evaluation by grouping, sorting by start_time, and extracting specified columns.

    Args:
        df: DataFrame containing trace data
        group_by_col: Column to group traces by (default: "context.trace_id")
        extract_cols: Dict mapping {output_key: source_column} to extract from each row
                     Can contain multiple columns to extract
        additional_data: Dict of additional data to include with each trace (default: None)
        filter_empty: Whether to filter out empty values (default: True)

    Returns:
        DataFrame with processed trace data ready for evaluation
    """
    # Group by specified column
    grouped = df.groupby(group_by_col)

    # Prepare results list
    results = []

    for group_id, group in grouped:
        # Always sort by start_time to ensure correct order
        group = group.sort_values("start_time")

        # Initialize a dict to store extracted data
        trace_data = {group_by_col: group[group_by_col].iloc[0]}

        # Extract and process each requested column
        for output_key, source_col in extract_cols.items():
            ordered_extracts = []
            # Iterate through rows as dictionaries to handle column names with dots
            for i, (_, row_data) in enumerate(group.reset_index(drop=True).iterrows()):
                # Convert row to dictionary for easier access
                row_dict = row_data.to_dict()
                value = row_dict.get(source_col)
                if not filter_empty or (value is not None and value):
                    ordered_extracts.append({str(i + 1): value})
            trace_data[output_key] = ordered_extracts

        # Add any additional data
        if additional_data:
            trace_data.update(additional_data)

        # Add to results
        results.append(trace_data)

    # Convert to DataFrame
    return pd.DataFrame(results)


def extract_tool_calls(output_messages):
    if not output_messages:
        return []

    tool_calls = []
    for message in output_messages:
        if "message.tool_calls" in message:
            for tool_call in message["message.tool_calls"]:
                tool_calls.append(
                    {
                        "name": tool_call["tool_call.function.name"],
                        "arguments": tool_call["tool_call.function.arguments"],
                    }
                )
    return tool_calls

# Data prep - filter traces for each agent
eval_traces = filter_spans_by_trace_criteria(
    df=trajectory_df,
    trace_filters={
        "name": {"contains": "bull_agent|bear_agent|trading_strategy_orchestrator"}
    },
    span_filters={"attributes.openinference.span.kind": {"==": "LLM"}},
)

eval_traces["tool_calls"] = eval_traces["attributes.llm.output_messages"].apply(
    extract_tool_calls
)

# aggregate tool calls by trace id
tool_calls_df = prepare_trace_data_for_evaluation(
    df=eval_traces,
    extract_cols={
        "tool_calls": "tool_calls",
        "attributes.llm.tools": "attributes.llm.tools",
        "attributes.input.value": "attributes.input.value",
    },  # can also add any additional columns to the dataframe
    # additional_data={"reference_outputs": reference_outputs},
)

# Run evaluations on trace level - aggregated tool calls / trajectory
import nest_asyncio

nest_asyncio.apply()

evaluations_df = llm_classify(
    dataframe=tool_calls_df,
    template=TRAJECTORY_ACCURACY_PROMPT_WITHOUT_REFERENCE,
    model=model,
    rails=["correct", "incorrect"],
    provide_explanation=True,
    verbose=False,
    concurrency=5,
)

# Prep data for upload
# Copy evaluations and add trace_id from tool_calls_df (they're in same order)
eval_df = evaluations_df.copy()
eval_df["context.trace_id"] = tool_calls_df["context.trace_id"].values

# Get the root span_id for each trace_id
root_spans = trajectory_df[trajectory_df["parent_id"].isnull()][
    ["context.trace_id", "context.span_id"]
]

# Merge evaluations with root spans to get the span_id
eval_df = pd.merge(eval_df, root_spans, on="context.trace_id", how="left")

# Rename context.span_id to span_id for upload
eval_df = eval_df.rename(columns={"context.span_id": "span_id"})

# Create score column
eval_df["score"] = eval_df["label"].apply(lambda x: 1 if x == "correct" else 0)

# Clean up explanation
eval_df["explanation"] = eval_df["explanation"].str.replace(
    r"^(correct|incorrect)\s*\n*EXPLANATION\s*\n*", "", regex=True
)

# Select columns for upload
eval_df = eval_df[["span_id", "label", "score", "explanation"]]

print("\nFinal eval_df for upload:")
print(eval_df.head())

# Upload to Phoenix
client.spans.log_span_annotations_dataframe(
    dataframe=eval_df, annotation_name="agent_trajectory", annotator_kind="LLM"
)

delete_bear_agent = False
delete_bull_agent = False

if delete_bear_agent:
    client.agent_engines.delete(deployed_bear.api_resource.name, force=True)
if delete_bull_agent:
    client.agent_engines.delete(deployed_bull.api_resource.name, force=True)