from pathlib import Path



p = Path("/app/api/.venv/lib/python3.12/site-packages/dify_trace_langfuse/langfuse_trace.py")

t = p.read_text()

marker = 'Langfuse v4 requires observation body.id; auto-assign if missing'

if marker in t:

    print("already patched")

    raise SystemExit(0)



old_span = """        data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {}

        try:

            body = CreateSpanBody(

                id=data.get("id"),

"""

new_span = """        data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {}

        try:

            # Langfuse v4 requires observation body.id; auto-assign if missing

            if not data.get("id"):

                data["id"] = str(uuid.uuid4())

            body = CreateSpanBody(

                id=data.get("id"),

"""

old_gen = """            usage_data = data.pop("usage", None)

            usage = None

            if usage_data:

                usage = Usage(

                    input=usage_data.get("input", 0) or 0,

                    output=usage_data.get("output", 0) or 0,

                    total=usage_data.get("total", 0) or 0,

                    unit=usage_data.get("unit"),

                    input_cost=usage_data.get("inputCost"),

                    output_cost=usage_data.get("outputCost"),

                    total_cost=usage_data.get("totalCost"),

                )



            body = CreateGenerationBody(

                id=data.get("id"),

"""

new_gen = """            usage_data = data.pop("usage", None)

            usage = None

            if usage_data:

                usage = Usage(

                    input=usage_data.get("input", 0) or 0,

                    output=usage_data.get("output", 0) or 0,

                    total=usage_data.get("total", 0) or 0,

                    unit=usage_data.get("unit"),

                    input_cost=usage_data.get("inputCost"),

                    output_cost=usage_data.get("outputCost"),

                    total_cost=usage_data.get("totalCost"),

                )



            # Langfuse v4 requires observation body.id; auto-assign if missing

            if not data.get("id"):

                data["id"] = str(uuid.uuid4())

            body = CreateGenerationBody(

                id=data.get("id"),

"""

if old_span not in t:

    raise SystemExit("span block not found")

if old_gen not in t:

    raise SystemExit("gen block not found")

t = t.replace(old_span, new_span, 1).replace(old_gen, new_gen, 1)

p.write_text(t)

print("patched ok")