# Coalex SDK

A Python package for OpenTelemetry integration with Coalex.ai observability platform. Automatically capture token usage, latency, and traces from all your LLM applications.

## Installation

### Basic Installation

```bash
pip install coalex
```

### With Auto-Instrumentation (Recommended)

Install with support for common LLM frameworks:

```bash
pip install "coalex[auto-instrument]"
```

This includes instrumentation for: OpenAI, LlamaIndex, LangChain, VertexAI, Anthropic, and AWS Bedrock.

### Framework-Specific Installation

```bash
# OpenAI only
pip install "coalex[openai]"

# LangChain only
pip install "coalex[langchain]"

# LlamaIndex only
pip install "coalex[llamaindex]"

# VertexAI only
pip install "coalex[vertexai]"

# All supported frameworks
pip install "coalex[all]"
```

## Quick Start (Recommended)

The easiest way to get started is with auto-instrumentation, which automatically captures tokens and traces from all installed LLM libraries:

```python
import coalex
from openai import OpenAI

# One-time setup - register and auto-instrument
tracer_provider = coalex.register(agent_id="YOUR_AGENT_ID")
results = coalex.auto_instrument(tracer_provider=tracer_provider)

# See what was instrumented
instrumented = [k for k, v in results.items() if v == "success"]
print(f"Instrumented: {', '.join(instrumented)}")

# Use your LLM library normally - tokens are automatically captured!
client = OpenAI()

with coalex.coalex_context(
    request_id="req_001",
    prompt_version="v1.0.0"
):
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": "Hello!"}]
    )
    print(response.choices[0].message.content)
```

## Supported Frameworks

Coalex automatically instruments these frameworks via OpenInference:

**LLM Frameworks & Agents:**
- OpenAI SDK
- LangChain
- LlamaIndex
- DSPy
- CrewAI
- Haystack
- Instructor
- Guardrails

**LLM Providers:**
- Google VertexAI
- AWS Bedrock
- Anthropic Claude
- MistralAI
- Groq
- liteLLM

## Features

- **Auto-Instrumentation**: One call instruments all installed LLM libraries
- **Token Tracking**: Automatically captures input/output tokens for all LLM calls
- **Zero Code Changes**: Use your LLM libraries normally - tracing happens automatically
- **Resilient**: Never crashes your app - errors are logged but don't interrupt execution
- **Context Management**: Proper span hierarchy and attribute propagation with `coalex_context()`
- **OpenInference Standard**: Uses industry-standard semantic conventions
- **Authentication**: Automatic authentication using agent_id

## Usage

### Auto-Instrumentation (Recommended)

Auto-instrument all installed LLM libraries:

```python
import coalex

# Register and auto-instrument everything
tracer_provider = coalex.register(agent_id="YOUR_AGENT_ID")
results = coalex.auto_instrument(tracer_provider=tracer_provider)
```

### Selective Instrumentation

Only instrument specific libraries:

```python
# Only OpenAI and LangChain
results = coalex.auto_instrument(
    tracer_provider=tracer_provider,
    include_libraries=["openai", "langchain"]
)
```

Exclude certain libraries:

```python
# Everything except Bedrock and Groq
results = coalex.auto_instrument(
    tracer_provider=tracer_provider,
    exclude_libraries=["bedrock", "groq"]
)
```

### Individual Instrumentors

For fine-grained control:

```python
# Instrument specific libraries individually
coalex.instrument_openai(tracer_provider)
coalex.instrument_langchain(tracer_provider)
coalex.instrument_vertexai(tracer_provider)
```

### Context Management

Use `coalex_context()` to group related LLM calls and add metadata:

```python
with coalex.coalex_context(
    request_id="user_session_123",
    prompt_version="v2.1.0"
):
    # All LLM calls here are grouped under this context
    response1 = client.chat.completions.create(...)
    response2 = client.chat.completions.create(...)
```

## Configuration

### `register()` Parameters

```python
tracer_provider = coalex.register(
    agent_id="YOUR_AGENT_ID",              # Required: Your agent ID
    endpoint=None,                          # Optional: Custom OTLP endpoint
    additional_attributes={},               # Optional: Extra attributes for all spans
    suppress_export_errors=True,            # Optional: Suppress export errors
    suppress_registration_errors=True,      # Optional: Never crash on registration errors
)
```

### `auto_instrument()` Parameters

```python
results = coalex.auto_instrument(
    tracer_provider=tracer_provider,        # Optional: TracerProvider from register()
    suppress_errors=True,                   # Optional: Suppress instrumentation errors
    include_libraries=None,                 # Optional: Only instrument these libraries
    exclude_libraries=None,                 # Optional: Skip these libraries
)
```

## Examples

Complete working examples are available in the `examples/` directory:

- **`openai_example.py`**: OpenAI SDK with streaming and non-streaming
- **`langchain_example.py`**: LangChain with LCEL chains and batch processing
- **`langgraph_async_example.py`**: Async LangGraph with tool calling
- **`vertexai_example.py`**: Google VertexAI with streaming

To run an example:

```bash
# Set your API keys
export OPENAI_API_KEY="your-key"
export COALEX_AGENT_ID="your-agent-id"

# Install dependencies
pip install "coalex[openai]" openai

# Run the example
python examples/openai_example.py
```

## Troubleshooting

### Token Counts Showing as 0?

If you see token counts as 0 in your Coalex dashboard:

1. **Use auto-instrumentation**: The most common issue is forgetting to instrument your LLM libraries
   ```python
   coalex.auto_instrument(tracer_provider=tracer_provider)
   ```

2. **For streaming requests**: Some providers require special configuration:
   ```python
   # OpenAI streaming - include usage in stream
   stream = client.chat.completions.create(
       model="gpt-4",
       messages=[...],
       stream=True,
       stream_options={"include_usage": True}  # Important!
   )
   ```

3. **Check instrumentation results**: Verify your library was actually instrumented:
   ```python
   results = coalex.auto_instrument(tracer_provider=tracer_provider)
   print(results)  # Should show "success" for your libraries
   ```

### Registration Errors

Coalex SDK is designed to never crash your application. If registration fails:

- Errors are logged but don't raise exceptions (by default)
- A no-op tracer is used (your app continues without tracing)
- Check logs for details: `logger.warning("Failed to register Coalex tracing...")`

To raise exceptions on registration errors:

```python
tracer_provider = coalex.register(
    agent_id="YOUR_AGENT_ID",
    suppress_registration_errors=False  # Will raise on errors
)
```

## Migration Guide

### Upgrading from v0.4.x to v0.5.0

**Before (v0.4.x)** - Manual instrumentation:

```python
from coalex.otel import register, coalex_context
from openinference.instrumentation.openai import OpenAIInstrumentor
from openinference.instrumentation.langchain import LangChainInstrumentor

tracer_provider = register(agent_id="YOUR_AGENT_ID")

# Had to manually instrument each library
OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)
LangChainInstrumentor().instrument(tracer_provider=tracer_provider)
# ... and so on for each library
```

**After (v0.5.0)** - Auto-instrumentation:

```python
import coalex

tracer_provider = coalex.register(agent_id="YOUR_AGENT_ID")

# One line instruments everything!
coalex.auto_instrument(tracer_provider=tracer_provider)
```

**Benefits:**
- No need to import individual instrumentors
- Automatically instruments all installed libraries
- Easier to maintain - works with new libraries without code changes
- More reliable - gracefully handles missing libraries

**Backward Compatibility:**
Manual instrumentation still works! You can continue using the old approach if needed.

## Advanced Usage

### Custom Endpoint

Use a custom OTLP endpoint:

```python
tracer_provider = coalex.register(
    agent_id="YOUR_AGENT_ID",
    endpoint="https://your-custom-endpoint.com/v1/traces"
)
```

### Additional Resource Attributes

Add custom attributes to all spans:

```python
tracer_provider = coalex.register(
    agent_id="YOUR_AGENT_ID",
    service_name="my-chatbot",
    additional_attributes={
        "environment": "production",
        "version": "2.1.0",
        "team": "ai-platform"
    }
)
```

### Manual Span Creation

For advanced use cases, create custom spans:

```python
from opentelemetry import trace

tracer = trace.get_tracer(__name__)

with tracer.start_as_current_span("custom_operation") as span:
    span.set_attribute("custom.attribute", "value")
    # Your code here
```

## Contributing

Contributions are welcome! Please open an issue or pull request on GitHub.

## License

See LICENSE file for details.

