# llm-observatory

Drop-in LLM observability for Python. One import change, all calls logged automatically.

## Install

```bash
pip install llm-observatory
```

## Usage

```python
from llm_observatory import configure, OpenAI

configure(api_key="your-api-key")

client = OpenAI(api_key="sk-...")
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "Hello"}],
)
```

Supports `OpenAI`, `AsyncOpenAI`, `AzureOpenAI`, `AsyncAzureOpenAI`, `Anthropic`, and `AsyncAnthropic`.

You can also set the `LLM_OBSERVATORY_API_KEY` environment variable instead of passing `api_key` to `configure()` — the SDK picks it up automatically.

## Langfuse integration

Already using Langfuse? `wrap_langfuse()` captures **all** Langfuse traces, generations, and spans — including metadata like `user_id`, `session_id`, and `tags` — and forwards them to Context0 automatically. Zero code changes to your existing Langfuse setup.

Requires **Langfuse Python SDK v3.0.0+** (OTel-native).

### Setup

```python
from llm_observatory import configure, wrap_langfuse

configure(api_key="your-api-key")

# Call once after Langfuse is initialized — adds Context0 processor
wrap_langfuse()
```

### What gets captured

After calling `wrap_langfuse()`, every Langfuse observation is automatically forwarded to Context0:

```python
from langfuse.openai import OpenAI as LangfuseOpenAI
from langfuse import propagate_attributes

client = LangfuseOpenAI(api_key="sk-...")

# Per-request metadata (user_id, session_id, tags) — captured automatically
with propagate_attributes(user_id="alice", session_id="sess-1", tags=["prod"]):
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "What is observability?"}],
    )

# @observe decorator — captured automatically
from langfuse.decorators import observe

@observe(as_type="generation")
def my_llm_call(prompt):
    return client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
    )

# Manual traces with start_as_current_observation — captured automatically
from langfuse import Langfuse

langfuse = Langfuse()
with langfuse.start_as_current_observation(as_type="span", name="rag-pipeline"):
    # Nested spans create parent-child relationships in Context0
    with langfuse.start_as_current_observation(as_type="span", name="retrieval"):
        docs = vector_db.search(query)

    with langfuse.start_as_current_observation(as_type="generation", name="completion"):
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": f"{docs}\n\n{query}"}],
        )
```

Works with all Langfuse v3+ patterns: `@observe`, `LangfuseOpenAI`, `start_as_current_observation()`, `propagate_attributes()`, and framework integrations (LangChain, LlamaIndex).

### `wrap()` — per-client alternative

If you only want to observe a **specific** client instance (rather than all Langfuse observations globally), use `wrap()` instead:

```python
from langfuse.openai import OpenAI as LangfuseOpenAI
from llm_observatory import configure, wrap

configure(api_key="your-api-key")

client = wrap(LangfuseOpenAI(api_key="sk-..."))

# Both Langfuse AND Context0 capture this call
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "Hello"}],
)
```

`wrap()` works with any client that has a `.chat` (OpenAI-shaped) or `.messages` (Anthropic-shaped) interface. If the client is already a Context0 wrapper, it's returned as-is — no double-wrapping.

**When to use which:** Use `wrap_langfuse()` to capture everything Langfuse sees globally (recommended). Use `wrap()` if you only want Context0 on specific client instances.

## Tracing

Group related LLM calls into traces to see your full pipeline as a tree.

### `@observe` decorator

Automatically creates a trace (top-level) or span (nested). Captures function args as input and return value as output.

```python
import llm_observatory as obs

client = obs.OpenAI()

@obs.observe
def retrieve_docs(query: str):
    return vector_db.search(query)

@obs.observe
def rag_pipeline(question: str):
    docs = retrieve_docs(question)       # child span
    response = client.chat.completions.create(  # auto-captured as generation
        model="gpt-4o",
        messages=[{"role": "user", "content": f"{docs}\n\n{question}"}],
    )
    return response.choices[0].message.content
```

Options:
- `@obs.observe(name="custom-name")` — override the span name (default: function name)
- `@obs.observe(capture_input=False)` — don't log args (for sensitive data)
- `@obs.observe(capture_output=False)` — don't log return value

Works with both sync and async functions.

### Context managers

Use `trace()` and `span()` for inline code blocks that aren't standalone functions:

```python
with obs.trace(name="rag-pipeline", input={"query": q}) as t:
    with obs.span(name="vector-search") as s:
        results = search(q)
        s.set_output(results)

    response = client.chat.completions.create(model="gpt-4o", messages=[...])
    t.set_output(response.choices[0].message.content)
```

Mix freely — `@observe` and `with span()` compose within the same trace.

### Cross-service propagation

Pass trace context across services (e.g., API to SQS worker) via W3C traceparent:

```python
# Producer — get traceparent inside a trace
with obs.trace(name="api-request"):
    traceparent = obs.get_current_traceparent()
    sqs.send_message(Body=json.dumps({"traceparent": traceparent, ...}))

# Consumer — restore context
msg = json.loads(sqs_message["Body"])
with obs.trace(name="worker", traceparent=msg["traceparent"]):
    client.chat.completions.create(...)  # appears as child in same trace
```

Use `obs.emit_completed_span()` to backfill spans with explicit start/end times (e.g., queue wait duration).

## Serverless / short-lived environments

In AWS Lambda, Convex, or any environment where the process freezes or terminates after your handler returns, buffered events may be lost before the background flush fires. Call `flush()` before returning:

```python
from llm_observatory import configure, OpenAI, flush

configure(api_key="your-api-key")
client = OpenAI(api_key="sk-...")

def handler(event, context):
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": event["prompt"]}],
    )

    flush()  # blocks until all buffered events are sent — call before the process freezes

    return {"statusCode": 200, "body": response.choices[0].message.content}
```

The SDK also registers an `atexit` handler that flushes on normal process shutdown, but Lambda freezes bypass `atexit` — so an explicit `flush()` is required.

### Using with Langfuse in serverless

If you're using `wrap_langfuse()`, Langfuse also buffers spans and needs flushing. The OTel TracerProvider's `shutdown()` flushes all registered span processors — both Langfuse and Context0 — in one call:

```python
from opentelemetry import trace

def handler(event, context):
    # ... your LLM calls ...

    trace.get_tracer_provider().shutdown()  # flushes both Langfuse and Context0

    return {"statusCode": 200, "body": "..."}
```
