import asyncio
from litellm import Router
import litellm
from langfuse.decorators import observe
import os
from langfuse.decorators import langfuse_context

# Configuration
MODEL_NAME = "your-model-name"  # Change to your deployment name
API_BASE = "https://your-endpoint.openai.azure.com"  # Insert your api base
API_VERSION = "2023-12-01-preview"
API_KEY = os.getenv("AZURE_API_KEY")
REDIS_URL = "redis://localhost:6379"

# Langfuse configuration
os.environ["LANGFUSE_HOST"] = "your-langfuse-host"
os.environ["LANGFUSE_PUBLIC_KEY"] = "your-public-key"
os.environ["LANGFUSE_SECRET_KEY"] = "your-secret-key"

# Configure LiteLLM callbacks
litellm.success_callback = ["langfuse"]
litellm.failure_callback = ["langfuse"]

# Initialize router
router = Router(
    model_list=[
        {
            "model_name": MODEL_NAME,
            "litellm_params": {
                "model": f"azure/{MODEL_NAME}",
                "api_base": API_BASE,
                "api_key": API_KEY,
                "api_version": API_VERSION,
            },
        }
    ],
    default_litellm_params={"acompletion": True},
    # Once REDIS is enabled here, langfuse integration sends the wrong
    # trace_id in parallel calls:
    redis_url=REDIS_URL,
)


async def call_llm(prompt: str):
    # Correct trace_id is printed here:
    print(
        "get_current_trace_id:",
        langfuse_context.get_current_trace_id(),
    )

    # Surprisingly, acompletion() works good, but we need
    # completions.create() to be fixed, as we need it for integration with
    # Instructor.
    # response = await router.acompletion(

    response = await router.chat.completions.create(
        model=MODEL_NAME,
        messages=[{"role": "user", "content": prompt}],
        metadata={
            "trace_id": langfuse_context.get_current_trace_id(),
            "generation_name": prompt,
            "debug_langfuse": True,
        },
    )
    return response


@observe()
async def process():
    # First call with Request1
    await call_llm("Tell me the result of 2+2")

    # Second call with Request2
    await call_llm("Do you like Math, yes or no?")


async def main():
    # Run two process functions in parallel
    await asyncio.gather(process(), process())


if __name__ == "__main__":
    asyncio.run(main())

get_current_trace_id: c45394a2-4fa0-4599-aa3c-88a101b35868
get_current_trace_id: fcb74aee-2de0-465e-b1f3-afd4730fe193
Real sent trace_id: fcb74aee-2de0-465e-b1f3-afd4730fe193
get_current_trace_id: c45394a2-4fa0-4599-aa3c-88a101b35868
Real sent trace_id: fcb74aee-2de0-465e-b1f3-afd4730fe193
get_current_trace_id: fcb74aee-2de0-465e-b1f3-afd4730fe193
Real sent trace_id: c45394a2-4fa0-4599-aa3c-88a101b35868
Real sent trace_id: fcb74aee-2de0-465e-b1f3-afd4730fe193

get_current_trace_id: 3a0d9972-9730-465e-9a63-840e9c8f8fd3
get_current_trace_id: 94e7c707-0bd7-47a9-8e25-bc8f8eca2b6d
Real sent trace_id: 94e7c707-0bd7-47a9-8e25-bc8f8eca2b6d
get_current_trace_id: 94e7c707-0bd7-47a9-8e25-bc8f8eca2b6d
Real sent trace_id: 3a0d9972-9730-465e-9a63-840e9c8f8fd3
get_current_trace_id: 3a0d9972-9730-465e-9a63-840e9c8f8fd3
Real sent trace_id: 94e7c707-0bd7-47a9-8e25-bc8f8eca2b6d
Real sent trace_id: 3a0d9972-9730-465e-9a63-840e9c8f8fd3