from brokle import Brokle
from brokle.evaluation import LLMJudge, evaluate
from dataclasses import dataclass
from typing import Optional

brokle = Brokle()

@dataclass
class QualityScores:
    accuracy: float
    helpfulness: float
    clarity: float
    conciseness: float
    tone: float
    overall: float
    details: dict

# Define dimension-specific judges
judges = {
    "accuracy": LLMJudge(
        name="accuracy",
        prompt="""Rate the factual accuracy of this response.

Question: {input}
Response: {output}
Reference (if available): {reference}

Consider:
- Are facts correct?
- Are any claims verifiable?
- Are there any errors or misconceptions?

Rate 0-1 where 1 is completely accurate.""",
        model="gpt-4o"
    ),

    "helpfulness": LLMJudge(
        name="helpfulness",
        prompt="""Rate how helpful this response is for the user's need.

User Query: {input}
Response: {output}

Consider:
- Does it address the actual question?
- Does it provide actionable information?
- Would the user get what they need?

Rate 0-1 where 1 is extremely helpful.""",
        model="gpt-4o-mini"
    ),

    "clarity": LLMJudge(
        name="clarity",
        prompt="""Rate the clarity and structure of this response.

Response: {output}

Consider:
- Is it well-organized?
- Is the language clear and unambiguous?
- Is it easy to follow?

Rate 0-1 where 1 is crystal clear.""",
        model="gpt-4o-mini"
    ),

    "conciseness": LLMJudge(
        name="conciseness",
        prompt="""Rate the conciseness of this response.

Query: {input}
Response: {output}

Consider:
- Is it appropriately detailed for the question?
- Is there unnecessary repetition?
- Could it be shorter without losing value?

Rate 0-1 where 1 is optimally concise (not too short, not too long).""",
        model="gpt-4o-mini"
    ),

    "tone": LLMJudge(
        name="tone",
        prompt="""Rate the appropriateness of the tone.

Context: {context}
Response: {output}

Consider:
- Is the tone professional and appropriate?
- Does it match the expected context?
- Is it respectful and neutral?

Rate 0-1 where 1 is perfectly appropriate tone.""",
        model="gpt-4o-mini"
    )
}

def evaluate_quality(
    query: str,
    response: str,
    reference: str = None,
    context: str = "general assistant",
    weights: dict = None
) -> QualityScores:
    """Evaluate response quality across multiple dimensions."""

    # Default weights
    if weights is None:
        weights = {
            "accuracy": 0.3,
            "helpfulness": 0.3,
            "clarity": 0.2,
            "conciseness": 0.1,
            "tone": 0.1
        }

    with brokle.start_as_current_span(
        name="quality_evaluation",
        metadata={"dimensions": list(weights.keys())}
    ) as span:
        scores = {}
        details = {}

        # Evaluate each dimension
        for dimension, judge in judges.items():
            result = judge.evaluate(
                input=query,
                output=response,
                reference=reference or "Not provided",
                context=context
            )

            scores[dimension] = result.score
            details[dimension] = result.comment

            # Record individual scores
            span.score(
                name=dimension,
                value=result.score,
                comment=result.comment[:200] if result.comment else None
            )

        # Calculate weighted overall score
        overall = sum(
            scores[dim] * weights[dim]
            for dim in weights
        )

        span.score(name="overall_quality", value=overall)
        span.update(
            output={
                "overall": overall,
                "dimensions": scores
            },
            metadata={"overall_quality": overall}
        )

        return QualityScores(
            accuracy=scores["accuracy"],
            helpfulness=scores["helpfulness"],
            clarity=scores["clarity"],
            conciseness=scores["conciseness"],
            tone=scores["tone"],
            overall=overall,
            details=details
        )

# Usage
query = "How do I reset my password?"
response = """
To reset your password, follow these steps:

1. Go to the login page
2. Click "Forgot Password"
3. Enter your email address
4. Check your inbox for a reset link
5. Click the link and create a new password

Make sure your new password is at least 8 characters with a mix of letters and numbers.
"""

quality = evaluate_quality(query, response)
print(f"Overall Quality: {quality.overall:.0%}")
print(f"  Accuracy: {quality.accuracy:.0%}")
print(f"  Helpfulness: {quality.helpfulness:.0%}")
print(f"  Clarity: {quality.clarity:.0%}")
print(f"  Conciseness: {quality.conciseness:.0%}")
print(f"  Tone: {quality.tone:.0%}")

brokle.flush()

import { Brokle, LLMJudge } from 'brokle';

const brokle = new Brokle();

const judges = {
  accuracy: new LLMJudge({
    name: 'accuracy',
    prompt: `Rate the factual accuracy. Question: {input} Response: {output}
Rate 0-1 where 1 is completely accurate.`,
    model: 'gpt-4o',
  }),
  helpfulness: new LLMJudge({
    name: 'helpfulness',
    prompt: `Rate how helpful this is. Query: {input} Response: {output}
Rate 0-1 where 1 is extremely helpful.`,
    model: 'gpt-4o-mini',
  }),
  clarity: new LLMJudge({
    name: 'clarity',
    prompt: `Rate the clarity. Response: {output}
Rate 0-1 where 1 is crystal clear.`,
    model: 'gpt-4o-mini',
  }),
};

async function evaluateQuality(query, response, weights = null) {
  weights = weights || { accuracy: 0.4, helpfulness: 0.4, clarity: 0.2 };

  return brokle.withSpan({ name: 'quality_evaluation' }, async (span) => {
    const scores = {};

    for (const [dimension, judge] of Object.entries(judges)) {
      const result = await judge.evaluate({ input: query, output: response });
      scores[dimension] = result.score;
      span.score({ name: dimension, value: result.score });
    }

    const overall = Object.entries(weights).reduce(
      (sum, [dim, weight]) => sum + (scores[dim] || 0) * weight,
      0
    );

    span.score({ name: 'overall_quality', value: overall });

    return { ...scores, overall };
  });
}