from brokle import Brokle, wrap_openai
import openai
from dataclasses import dataclass
from typing import Optional

brokle = Brokle()
openai_client = wrap_openai(openai.OpenAI(), brokle=brokle)

@dataclass
class BatchResult:
    item: str
    output: Optional[str]
    success: bool
    error: Optional[str] = None
    tokens: int = 0

def process_item(item: str, index: int) -> BatchResult:
    """Process single item with error handling."""

    with brokle.start_as_current_generation(
        name="batch_item",
        model="gpt-4o-mini",
        metadata={"index": index, "item_preview": item[:50]}
    ) as gen:
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[{"role": "user", "content": f"Classify: {item}"}],
                max_tokens=100
            )

            output = response.choices[0].message.content
            tokens = response.usage.total_tokens

            gen.update(
                output=output,
                usage={
                    "input_tokens": response.usage.prompt_tokens,
                    "output_tokens": response.usage.completion_tokens
                }
            )

            return BatchResult(
                item=item,
                output=output,
                success=True,
                tokens=tokens
            )

        except Exception as e:
            gen.update(
                output=str(e),
                metadata={"error": True, "error_type": type(e).__name__}
            )
            gen.score(name="error", value=1, comment=str(e))

            return BatchResult(
                item=item,
                output=None,
                success=False,
                error=str(e)
            )

def process_batch(
    items: list[str],
    batch_name: str = "batch_job"
) -> dict:
    """Process batch with aggregated metrics."""

    with brokle.start_as_current_span(
        name=batch_name,
        metadata={
            "batch_size": len(items),
            "operation": "classification"
        }
    ) as span:
        results = []
        total_tokens = 0
        success_count = 0

        for i, item in enumerate(items):
            result = process_item(item, i)
            results.append(result)

            if result.success:
                success_count += 1
                total_tokens += result.tokens

            # Update progress periodically
            if (i + 1) % 10 == 0:
                span.set_attribute("progress", f"{i + 1}/{len(items)}")

        # Final metrics
        success_rate = success_count / len(items)
        span.update(
            output=f"Processed {len(items)} items",
            metadata={
                "success_count": success_count,
                "failure_count": len(items) - success_count,
                "success_rate": success_rate,
                "total_tokens": total_tokens,
                "avg_tokens_per_item": total_tokens / max(success_count, 1)
            }
        )

        # Score batch quality
        span.score(
            name="batch_success_rate",
            value=success_rate,
            comment=f"{success_count}/{len(items)} successful"
        )

        return {
            "results": results,
            "metrics": {
                "total": len(items),
                "success": success_count,
                "failed": len(items) - success_count,
                "success_rate": success_rate,
                "total_tokens": total_tokens
            }
        }

# Usage
items = [f"Product description {i}" for i in range(100)]
batch_output = process_batch(items, batch_name="product_classification")
print(f"Success rate: {batch_output['metrics']['success_rate']:.1%}")

brokle.flush()

import { Brokle } from 'brokle';
import { wrapOpenAI } from 'brokle-openai';
import OpenAI from 'openai';

const brokle = new Brokle();
const openai = wrapOpenAI(new OpenAI(), { brokle });

async function processItem(item, index) {
  return brokle.withGeneration(
    {
      name: 'batch_item',
      model: 'gpt-4o-mini',
      metadata: { index, itemPreview: item.slice(0, 50) },
    },
    async (gen) => {
      try {
        const response = await openai.chat.completions.create({
          model: 'gpt-4o-mini',
          messages: [{ role: 'user', content: `Classify: ${item}` }],
          max_tokens: 100,
        });

        const output = response.choices[0].message.content;
        const tokens = response.usage.total_tokens;

        gen.update({
          output,
          usage: {
            inputTokens: response.usage.prompt_tokens,
            outputTokens: response.usage.completion_tokens,
          },
        });

        return { item, output, success: true, tokens };
      } catch (error) {
        gen.update({
          output: error.message,
          metadata: { error: true, errorType: error.name },
        });
        gen.score({ name: 'error', value: 1, comment: error.message });

        return { item, output: null, success: false, error: error.message };
      }
    }
  );
}

async function processBatch(items, batchName = 'batch_job') {
  return brokle.withSpan(
    { name: batchName, metadata: { batchSize: items.length } },
    async (span) => {
      const results = [];
      let totalTokens = 0;
      let successCount = 0;

      for (let i = 0; i < items.length; i++) {
        const result = await processItem(items[i], i);
        results.push(result);

        if (result.success) {
          successCount++;
          totalTokens += result.tokens;
        }
      }

      const successRate = successCount / items.length;
      span.update({
        output: `Processed ${items.length} items`,
        metadata: {
          successCount,
          failureCount: items.length - successCount,
          successRate,
          totalTokens,
        },
      });

      span.score({
        name: 'batch_success_rate',
        value: successRate,
        comment: `${successCount}/${items.length} successful`,
      });

      return { results, metrics: { total: items.length, success: successCount } };
    }
  );
}