from brokle import Brokle

client = Brokle(api_key="bk_...")

# Using context manager (recommended)
with client.start_as_current_span(name="process_document") as span:
    result = process(document)
    span.update(output=result)

# Using decorator
from brokle import observe

@observe(name="analyze_text")
def analyze_text(text: str) -> dict:
    return {"sentiment": "positive", "confidence": 0.95}

import { Brokle } from 'brokle';

const client = new Brokle({ apiKey: 'bk_...' });

// Start and end span explicitly
const span = client.startSpan({ name: 'process_document' });

try {
  const result = await process(document);
  span.end({ output: result });
} catch (error) {
  span.end({ error });
  throw error;
}

with client.start_as_current_generation(
    name="summarize",
    model="gpt-4",
    input={"text": document_text}
) as gen:
    response = openai.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "Summarize the following text."},
            {"role": "user", "content": document_text}
        ]
    )

    gen.update(
        output=response.choices[0].message.content,
        usage={
            "prompt_tokens": response.usage.prompt_tokens,
            "completion_tokens": response.usage.completion_tokens
        },
        metadata={
            "temperature": 0.7,
            "max_tokens": 500
        }
    )

const gen = client.startGeneration({
  name: 'summarize',
  model: 'gpt-4',
  input: { text: documentText }
});

const response = await openai.chat.completions.create({
  model: 'gpt-4',
  messages: [
    { role: 'system', content: 'Summarize the following text.' },
    { role: 'user', content: documentText }
  ]
});

gen.end({
  output: response.choices[0].message.content,
  usage: {
    promptTokens: response.usage.prompt_tokens,
    completionTokens: response.usage.completion_tokens
  },
  attributes: {
    temperature: 0.7,
    maxTokens: 500
  }
});

with client.start_as_current_span(name="rag_pipeline") as parent:
    # Child span 1: Query embedding
    with client.start_as_current_span(name="embed_query") as embed_span:
        embedding = embed(query)
        embed_span.update(
            output={"dimensions": len(embedding)},
            metadata={"model": "text-embedding-3-small"}
        )

    # Child span 2: Vector search
    with client.start_as_current_span(
        name="vector_search",
        as_type="retrieval"
    ) as search_span:
        search_span.set_attribute("index", "products")
        search_span.set_attribute("top_k", 5)

        results = vector_store.search(embedding, top_k=5)

        search_span.update(output={
            "count": len(results),
            "scores": [r.score for r in results]
        })

    # Child span 3: LLM generation
    with client.start_as_current_generation(
        name="generate_answer",
        model="gpt-4"
    ) as gen_span:
        answer = generate(query, results)
        gen_span.update(output=answer)

    parent.update(output=answer)

const parent = client.startSpan({ name: 'rag_pipeline' });

// Child span 1: Query embedding
const embedSpan = client.startSpan({
  name: 'embed_query',
  parentSpanId: parent.spanId
});
const embedding = await embed(query);
embedSpan.end({
  output: { dimensions: embedding.length },
  attributes: { model: 'text-embedding-3-small' }
});

// Child span 2: Vector search
const searchSpan = client.startSpan({
  name: 'vector_search',
  type: 'retrieval',
  parentSpanId: parent.spanId,
  attributes: { index: 'products', topK: 5 }
});
const results = await vectorStore.search(embedding, 5);
searchSpan.end({
  output: {
    count: results.length,
    scores: results.map(r => r.score)
  }
});

// Child span 3: LLM generation
const genSpan = client.startGeneration({
  name: 'generate_answer',
  model: 'gpt-4',
  parentSpanId: parent.spanId
});
const answer = await generate(query, results);
genSpan.end({ output: answer });

parent.end({ output: answer });

with client.start_as_current_span(name="process_order") as span:
    # Set individual attributes
    span.set_attribute("order_id", order.id)
    span.set_attribute("customer_id", customer.id)
    span.set_attribute("item_count", len(order.items))
    span.set_attribute("total_value", order.total)
    span.set_attribute("priority", "high")

    # Or via metadata in update()
    span.update(
        metadata={
            "payment_method": order.payment_method,
            "shipping_tier": order.shipping_tier
        }
    )

    result = process_order(order)
    span.update(output=result)

const span = client.startSpan({
  name: 'process_order',
  attributes: {
    orderId: order.id,
    customerId: customer.id,
    itemCount: order.items.length,
    totalValue: order.total,
    priority: 'high'
  }
});

// Add more attributes later
span.setAttributes({
  paymentMethod: order.paymentMethod,
  shippingTier: order.shippingTier
});

const result = await processOrder(order);
span.end({ output: result });

with client.start_as_current_span(name="multi_step_process") as span:
    # Update as you progress
    span.set_attribute("step", "preprocessing")
    preprocessed = preprocess(data)

    span.set_attribute("step", "analysis")
    analysis = analyze(preprocessed)

    span.set_attribute("step", "postprocessing")
    result = postprocess(analysis)

    # Final update with output
    span.update(
        output=result,
        metadata={
            "preprocessing_time_ms": 100,
            "analysis_time_ms": 500,
            "postprocessing_time_ms": 50
        }
    )

const span = client.startSpan({ name: 'multi_step_process' });

span.setAttributes({ step: 'preprocessing' });
const preprocessed = await preprocess(data);

span.setAttributes({ step: 'analysis' });
const analysis = await analyze(preprocessed);

span.setAttributes({ step: 'postprocessing' });
const result = await postprocess(analysis);

span.end({
  output: result,
  attributes: {
    preprocessingTimeMs: 100,
    analysisTimeMs: 500,
    postprocessingTimeMs: 50
  }
});

with client.start_as_current_span(name="user_interaction") as span:
    # Link to user and session
    span.update_trace(
        user_id="user_123",
        session_id="session_456"
    )

    # Set trace-level tags
    span.set_attribute("environment", "production")
    span.set_attribute("version", "2.1.0")

    response = handle_interaction(request)
    span.update(output=response)

const span = client.startSpan({
  name: 'user_interaction',
  attributes: {
    userId: 'user_123',
    sessionId: 'session_456',
    environment: 'production',
    version: '2.1.0'
  }
});

const response = await handleInteraction(request);
span.end({ output: response });

with client.start_as_current_span(name="risky_operation") as span:
    try:
        result = perform_operation()
        span.update(output=result)
    except ValidationError as e:
        span.update(
            error=str(e),
            metadata={"error_type": "validation", "field": e.field}
        )
        raise
    except ExternalAPIError as e:
        span.update(
            error=str(e),
            metadata={
                "error_type": "external_api",
                "status_code": e.status_code,
                "retryable": e.retryable
            }
        )
        raise
    except Exception as e:
        span.update(error=f"Unexpected error: {str(e)}")
        raise

const span = client.startSpan({ name: 'risky_operation' });

try {
  const result = await performOperation();
  span.end({ output: result });
} catch (error) {
  if (error instanceof ValidationError) {
    span.end({
      error: error.message,
      attributes: { errorType: 'validation', field: error.field }
    });
  } else if (error instanceof ExternalAPIError) {
    span.end({
      error: error.message,
      attributes: {
        errorType: 'external_api',
        statusCode: error.statusCode,
        retryable: error.retryable
      }
    });
  } else {
    span.end({ error: `Unexpected error: ${error.message}` });
  }
  throw error;
}

with client.start_as_current_generation(
    name="streaming_response",
    model="gpt-4"
) as gen:
    stream = openai.chat.completions.create(
        model="gpt-4",
        messages=messages,
        stream=True
    )

    collected_content = []
    first_token_time = None

    for chunk in stream:
        if chunk.choices[0].delta.content:
            if first_token_time is None:
                first_token_time = time.time()
            collected_content.append(chunk.choices[0].delta.content)
            yield chunk.choices[0].delta.content

    full_response = "".join(collected_content)

    gen.update(
        output=full_response,
        metadata={
            "streaming": True,
            "time_to_first_token_ms": first_token_time * 1000
        }
    )

const gen = client.startGeneration({
  name: 'streaming_response',
  model: 'gpt-4'
});

const stream = await openai.chat.completions.create({
  model: 'gpt-4',
  messages,
  stream: true
});

const chunks = [];
let firstTokenTime = null;

for await (const chunk of stream) {
  if (chunk.choices[0]?.delta?.content) {
    if (firstTokenTime === null) {
      firstTokenTime = Date.now();
    }
    chunks.push(chunk.choices[0].delta.content);
    yield chunk.choices[0].delta.content;
  }
}

const fullResponse = chunks.join('');

gen.end({
  output: fullResponse,
  attributes: {
    streaming: true,
    timeToFirstTokenMs: firstTokenTime
  }
});