# Using OpenTelemetry for distributed tracing
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

# Configure the tracer
resource = Resource(attributes={SERVICE_NAME: "ai-agent-service"})
tracer_provider = TracerProvider(resource=resource)
otlp_exporter = OTLPSpanExporter(endpoint="otel-collector:4317", insecure=True)
span_processor = BatchSpanProcessor(otlp_exporter)
tracer_provider.add_span_processor(span_processor)
trace.set_tracer_provider(tracer_provider)

# Get a tracer
tracer = trace.get_tracer(__name__)

# Example of tracing a function
def process_user_query(query, user_id):
    with tracer.start_as_current_span("process_user_query") as span:
        # Add relevant attributes to the span
        span.set_attribute("user.id", user_id)
        span.set_attribute("query.length", len(query))
        
        # Trace the preprocessing step
        with tracer.start_as_current_span("preprocess_query"):
            processed_query = preprocess(query)
            
        # Trace the agent execution
        with tracer.start_as_current_span("execute_agent"):
            response = agent.run(processed_query)
            span.set_attribute("response.length", len(response))
            
        # Trace the postprocessing step
        with tracer.start_as_current_span("postprocess_response"):
            final_response = postprocess(response)
            
        return final_response

# Using a middleware approach to track token usage
import time
import tiktoken
from prometheus_client import Counter, Histogram

# Set up metrics
TOKENS_COUNTER = Counter('llm_tokens_total', 'Total tokens used', ['model', 'type'])
TOKENS_COST_COUNTER = Counter('llm_tokens_cost_total', 'Total cost of tokens used', ['model'])
LATENCY_HISTOGRAM = Histogram('llm_request_duration_seconds', 'Time spent processing LLM requests', ['model'])

# Initialize tokenizer
encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")

class TokenUsageMiddleware:
    def __init__(self, llm_client):
        self.llm_client = llm_client
        self.price_per_1k_tokens = {
            "gpt-3.5-turbo": {"input": 0.0015, "output": 0.002},
            "gpt-4": {"input": 0.03, "output": 0.06},
            # Add other models as needed
        }
    
    def count_tokens(self, text):
        """Count the number of tokens in the text"""
        tokens = encoding.encode(text)
        return len(tokens)
    
    def calculate_cost(self, model, input_tokens, output_tokens):
        """Calculate the cost based on token usage"""
        if model not in self.price_per_1k_tokens:
            return 0
        
        input_cost = (input_tokens / 1000) * self.price_per_1k_tokens[model]["input"]
        output_cost = (output_tokens / 1000) * self.price_per_1k_tokens[model]["output"]
        return input_cost + output_cost
    
    def completion(self, prompt, model="gpt-3.5-turbo", **kwargs):
        """Wrapper around the LLM completion call with token tracking"""
        input_tokens = self.count_tokens(prompt)
        TOKENS_COUNTER.labels(model=model, type="input").inc(input_tokens)
        
        start_time = time.time()
        response = self.llm_client.completion(prompt, model=model, **kwargs)
        duration = time.time() - start_time
        LATENCY_HISTOGRAM.labels(model=model).observe(duration)
        
        output_tokens = self.count_tokens(response.text)
        TOKENS_COUNTER.labels(model=model, type="output").inc(output_tokens)
        
        cost = self.calculate_cost(model, input_tokens, output_tokens)
        TOKENS_COST_COUNTER.labels(model=model).inc(cost)
        
        # Add metadata to response
        response.metadata = {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
            "cost": cost,
            "duration": duration
        }
        
        return response

# A simple evaluator for response quality
import json
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI

class ResponseEvaluator:
    def __init__(self):
        self.evaluation_llm = OpenAI(model_name="gpt-4", temperature=0)
        self.evaluation_prompt = PromptTemplate(
            input_variables=["query", "response", "criteria"],
            template="""Evaluate the quality of the following AI assistant response based on the given criteria.

USER QUERY: {query}

AI RESPONSE: {response}

EVALUATION CRITERIA:
{criteria}

Provide a score from 1-10 for each criterion and a brief explanation. Then give an overall score.
Return your evaluation as a JSON object with the following structure:
{{
  "criteria_scores": {{
    "criterion_name": {{
      "score": numeric_score,
      "explanation": "explanation"
    }},
    ...
  }},
  "overall_score": numeric_score,
  "summary": "brief summary of strengths and weaknesses"
}}"""
        )
        
        # Default evaluation criteria
        self.default_criteria = """
1. Accuracy: Does the response correctly answer the question or fulfill the request?
2. Completeness: Does the response address all aspects of the query?
3. Relevance: Is the response directly related to the query without unnecessary information?
4. Clarity: Is the response clear, well-organized, and easy to understand?
5. Helpfulness: Does the response provide practical value to the user?
"""

    def evaluate(self, query, response, criteria=None):
        """Evaluate the quality of an agent response"""
        eval_criteria = criteria if criteria else self.default_criteria
        
        prompt = self.evaluation_prompt.format(
            query=query,
            response=response,
            criteria=eval_criteria
        )
        
        result = self.evaluation_llm(prompt)
        
        try:
            evaluation = json.loads(result)
            return evaluation
        except json.JSONDecodeError:
            # Fallback if evaluation LLM doesn't return valid JSON
            return {
                "error": "Failed to parse evaluation",
                "raw_evaluation": result,
                "overall_score": 0
            }

// Frontend implementation of user feedback collection
const FeedbackComponent = () => {
  const [feedback, setFeedback] = useState(null);
  const [comment, setComment] = useState('');
  const [submitted, setSubmitted] = useState(false);
  const [conversationId, setConversationId] = useState(props.conversationId);

  const submitFeedback = async () => {
    try {
      const response = await fetch('/api/feedback', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          conversationId,
          rating: feedback,
          comment,
          timestamp: new Date().toISOString(),
        }),
      });
      
      if (response.ok) {
        setSubmitted(true);
      }
    } catch (error) {
      console.error('Error submitting feedback:', error);
    }
  };

  if (submitted) {
    return 
Thank you for your feedback!
;
  }

  return (
    <div className="feedback-container">
      <h4>Was this response helpful?</h4>
      <div className="rating-buttons">
        <button 
          className={`rating-button ${feedback === 'positive' ? 'active' : ''}`}
          onClick={() => setFeedback('positive')}
        >
          👍 Yes
        </button>
        <button 
          className={`rating-button ${feedback === 'negative' ? 'active' : ''}`}
          onClick={() => setFeedback('negative')}
        >
          👎 No
        </button>
      </div>
      
      {feedback && (
        <>
          <textarea 
            placeholder="Tell us more about your experience (optional)"
            value={comment}
            onChange={(e) => setComment(e.target.value)}
            rows={3}
          />
          <button className="submit-button" onClick={submitFeedback}>
            Submit Feedback
          </button>
        </>
      )}
    </div>
  );
};

# Prometheus Alerting Rules Example
groups:
- name: ai_agent_alerts
  rules:
  - alert: HighErrorRate
    expr: rate(ai_agent_errors_total[5m]) / rate(ai_agent_requests_total[5m]) > 0.05
    for: 2m
    labels:
      severity: critical
    annotations:
      summary: "High error rate detected"
      description: "Error rate is above 5% for the past 2 minutes"
      
  - alert: SlowResponseTime
    expr: histogram_quantile(0.95, rate(ai_agent_response_time_seconds_bucket[5m])) > 5
    for: 5m
    labels:
      severity: warning
    annotations:
      summary: "Slow response times detected"
      description: "95th percentile response time is above 5 seconds for the past 5 minutes"
      
  - alert: HighTokenUsage
    expr: increase(ai_agent_token_usage_total[1h]) > 1000000
    labels:
      severity: warning
    annotations:
      summary: "High token usage detected"
      description: "More than 1M tokens used in the past hour"
      
  - alert: LowUserSatisfaction
    expr: avg_over_time(ai_agent_user_satisfaction[1d]) < 3.5
    for: 6h
    labels:
      severity: warning
    annotations:
      summary: "Low user satisfaction scores"
      description: "Average user satisfaction below 3.5/5 for the past 6 hours"
      
  - alert: HighHallucinationRate
      expr: avg(ai_agent_hallucination_rate) > 0.08
      for: 30m
      labels:
        severity: critical
      annotations:
        summary: "High hallucination rate detected"
        description: "Agent hallucination rate above 8% for the past 30 minutes"

# Comprehensive agent evaluation framework
import json
import pandas as pd
from typing import List, Dict, Any, Optional
from dataclasses import dataclass

@dataclass
class EvaluationResult:
    category: str
    metric_name: str
    score: float
    benchmark: float
    pass_fail: bool
    details: Optional[Dict[str, Any]] = None

class AgentEvaluator:
    def __init__(self, benchmarks_path: str):
        # Load benchmarks from configuration
        with open(benchmarks_path, 'r') as f:
            self.benchmarks = json.load(f)
        
        # Initialize evaluators for different categories
        self.evaluators = {
            "technical": self._evaluate_technical_performance,
            "task_completion": self._evaluate_task_completion,
            "output_quality": self._evaluate_output_quality,
            "safety": self._evaluate_safety_compliance,
            "business_impact": self._evaluate_business_impact
        }
    
    def evaluate_agent(self, agent_id: str, evaluation_data: Dict[str, Any]) -> List[EvaluationResult]:
        """Run comprehensive evaluation of an agent"""
        results = []
        
        # Run evaluations for each category
        for category, evaluator_func in self.evaluators.items():
            category_results = evaluator_func(evaluation_data)
            results.extend(category_results)
        
        return results
    
    def _evaluate_technical_performance(self, data: Dict[str, Any]) -> List[EvaluationResult]:
        """Evaluate technical performance metrics"""
        results = []
        tech_metrics = self.benchmarks["technical"]
        
        # Latency evaluation
        avg_latency = data["technical"]["average_latency"]
        latency_benchmark = tech_metrics["latency"]["threshold"]
        results.append(EvaluationResult(
            category="technical",
            metric_name="average_latency",
            score=avg_latency,
            benchmark=latency_benchmark,
            pass_fail=avg_latency <= latency_benchmark,
            details={"unit": "seconds", "samples": data["technical"]["latency_samples"]}
        ))
        
        # Error rate evaluation
        error_rate = data["technical"]["error_rate"]
        error_benchmark = tech_metrics["error_rate"]["threshold"]
        results.append(EvaluationResult(
            category="technical",
            metric_name="error_rate",
            score=error_rate,
            benchmark=error_benchmark,
            pass_fail=error_rate <= error_benchmark,
            details={"total_requests": data["technical"]["total_requests"]}
        ))
        
        # Token efficiency
        tokens_per_task = data["technical"]["average_tokens_per_task"]
        token_benchmark = tech_metrics["tokens_per_task"]["threshold"]
        results.append(EvaluationResult(
            category="technical",
            metric_name="tokens_per_task",
            score=tokens_per_task,
            benchmark=token_benchmark,
            pass_fail=tokens_per_task <= token_benchmark
        ))
        
        return results
    
    def _evaluate_task_completion(self, data: Dict[str, Any]) -> List[EvaluationResult]:
        """Evaluate task completion metrics"""
        # Implementation would be similar to technical performance
        # but focused on task completion metrics
        return []
    
    def _evaluate_output_quality(self, data: Dict[str, Any]) -> List[EvaluationResult]:
        """Evaluate output quality metrics"""
        # Implementation for output quality metrics
        return []
    
    def _evaluate_safety_compliance(self, data: Dict[str, Any]) -> List[EvaluationResult]:
        """Evaluate safety compliance metrics"""
        # Implementation for safety metrics
        return []
    
    def _evaluate_business_impact(self, data: Dict[str, Any]) -> List[EvaluationResult]:
        """Evaluate business impact metrics"""
        # Implementation for business impact metrics
        return []
    
    def generate_evaluation_report(self, results: List[EvaluationResult]) -> Dict[str, Any]:
        """Generate a comprehensive evaluation report"""
        df = pd.DataFrame([
            {
                "category": r.category,
                "metric": r.metric_name,
                "score": r.score,
                "benchmark": r.benchmark,
                "status": "PASS" if r.pass_fail else "FAIL"
            }
            for r in results
        ])
        
        # Calculate summary statistics
        summary = {
            "total_metrics": len(results),
            "passed_metrics": sum(1 for r in results if r.pass_fail),
            "failed_metrics": sum(1 for r in results if not r.pass_fail),
            "pass_rate": sum(1 for r in results if r.pass_fail) / len(results) if results else 0,
            "categories": {
                category: {
                    "pass_rate": df[df["category"] == category]["status"].value_counts().get("PASS", 0) / len(df[df["category"] == category]) if len(df[df["category"] == category]) > 0 else 0
                }
                for category in df["category"].unique()
            }
        }
        
        return {
            "summary": summary,
            "details": [
                {
                    "category": r.category,
                    "metric": r.metric_name,
                    "score": r.score,
                    "benchmark": r.benchmark,
                    "status": "PASS" if r.pass_fail else "FAIL",
                    "details": r.details
                }
                for r in results
            ]
        }