import torch
import torch.nn as nn
import numpy as np
from collections import deque
import random

class QNetwork(nn.Module):
    def __init__(self, state_size, action_size, hidden_size=64):
        super().__init__()
        self.fc1 = nn.Linear(state_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, action_size)
        
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

class Agent:
    def __init__(self, state_size, action_size, lr=0.001):
        self.q_network = QNetwork(state_size, action_size)
        self.optimizer = torch.optim.Adam(self.q_network.parameters(), lr=lr)
        self.memory = deque(maxlen=10000)
        self.gamma = 0.95  # discount factor
        
    def act(self, state, epsilon=0.1):
        if random.random() < epsilon:
            return random.randint(0, self.q_network.fc3.out_features - 1)
        
        with torch.no_grad():
            state_t = torch.FloatTensor(state).unsqueeze(0)
            q_values = self.q_network(state_t)
            return q_values.argmax().item()
    
    def learn(self, batch_size=32):
        if len(self.memory) < batch_size:
            return
        
        batch = random.sample(self.memory, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        
        states = torch.FloatTensor(np.array(states))
        actions = torch.LongTensor(actions)
        rewards = torch.FloatTensor(rewards)
        next_states = torch.FloatTensor(np.array(next_states))
        dones = torch.FloatTensor(dones)
        
        # Q-learning update
        current_q = self.q_network(states).gather(1, actions.unsqueeze(1))
        next_q = self.q_network(next_states).max(1)[0].detach()
        target_q = rewards + (1 - dones) * self.gamma * next_q
        
        loss = nn.MSELoss()(current_q.squeeze(), target_q)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

from fastmcp import FastMCP
import azure.identity
from azure.identity import DefaultAzureCredential
import requests

mcp = FastMCP("Fabric Data Agent")

@mcp.tool()
def query_semantic_model(dataset_id: str, dax_query: str) -> dict:
    """Execute DAX query against Fabric semantic model.
    
    Args:
        dataset_id: The Power BI dataset/semantic model ID
        dax_query: DAX query string to execute
    
    Returns:
        Query results as dictionary
    """
    credential = DefaultAzureCredential()
    token = credential.get_token("https://analysis.windows.net/powerbi/api/.default")
    
    headers = {
        "Authorization": f"Bearer {token.token}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "queries": [{
            "query": dax_query
        }]
    }
    
    url = f"https://api.powerbi.com/v1.0/myorg/datasets/{dataset_id}/executeQueries"
    response = requests.post(url, headers=headers, json=payload)
    response.raise_for_status()
    
    return response.json()

if __name__ == "__main__":
    mcp.run()

import time
import functools
import random

def retry(errors=(Exception,), tries=4, base_delay=1, max_delay=30):
    """Retry on errors with exponential backoff and jitter.
    
    Args:
        errors: Exception types to retry on
        tries: Maximum number of attempts
        base_delay: Initial delay in seconds
        max_delay: Maximum delay cap in seconds
    """
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            for attempt in range(tries):
                try:
                    return fn(*args, **kwargs)
                except errors as e:
                    if attempt == tries - 1:
                        raise
                    
                    # Exponential backoff with jitter
                    delay = min(max_delay, base_delay * (2 ** attempt))
                    jittered_delay = delay * random.uniform(0.8, 1.2)
                    
                    print(f"⚠️  {type(e).__name__}: {e}")
                    print(f"   Retrying in {jittered_delay:.1f}s (attempt {attempt + 1}/{tries})")
                    time.sleep(jittered_delay)
            
        return wrapper
    return decorator

# Usage examples
@retry((ConnectionError, TimeoutError), tries=5)
def fetch_api(url):
    import requests
    return requests.get(url, timeout=10)

@retry((Exception,), tries=3, base_delay=2)
def process_data(data):
    # Your data processing logic
    pass