from obs.tools.base import BaseTool
from opentelemetry import trace


class LangfuseTool(BaseTool):

  def init(self):
    self.tracer = trace.get_tracer("langfuse")

  def push(self, trace_id, name, input, output, start_time, end_time, error):

    with self.tracer.start_as_current_span(name) as span:
      span.set_attribute("input", str(input))
      span.set_attribute("output", str(output))
      span.set_attribute("trace_id", trace_id)

      if error:
        span.set_attribute("error", error)