import palantyra
import uuid

# Generate session ID
session_id = str(uuid.uuid4())

# Method 1: Set at trace level
palantyra.set_trace_session_id(session_id)

# All subsequent calls in this trace belong to the session
response1 = client.chat.completions.create(...)
response2 = client.chat.completions.create(...)

# Method 2: Use with decorator
@palantyra.observe(session_id=session_id)
def chat_conversation():
    response1 = client.chat.completions.create(...)
    response2 = client.chat.completions.create(...)
    return [response1, response2]

import palantyra

# Track which user made the request
user_id = "user_abc123"
palantyra.set_trace_user_id(user_id)

# Now all traces include user information
response = client.chat.completions.create(...)

# Or use with context manager
with palantyra.start_as_current_span(
    name="User Chat",
    user_id=user_id,
    session_id=session_id
):
    response = client.chat.completions.create(...)

import palantyra
import uuid

def chat_session(user_id: str):
    session_id = str(uuid.uuid4())
    
    with palantyra.start_as_current_span(
        name="Chat Session",
        session_id=session_id,
        user_id=user_id,
        metadata={"conversation_type": "support"}
    ):
        messages = []
        
        # First message
        palantyra.set_span_attributes({"turn": 1})
        response1 = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": "Hello"}]
        )
        messages.append(response1.choices[0].message)
        
        # Second message
        palantyra.set_span_attributes({"turn": 2})
        response2 = client.chat.completions.create(
            model="gpt-4",
            messages=messages + [
                {"role": "user", "content": "How are you?"}
            ]
        )
        
        return response2

# Usage
result = chat_session("user_123")