import os
from langsmith import Client
from .base import BaseTool
from obs.utils.config import get_env


class LangSmithTool(BaseTool):

  def init(self):

    os.environ["LANGCHAIN_TRACING_V2"] = get_env("LANGCHAIN_TRACING_V2")
    os.environ["LANGCHAIN_API_KEY"] = get_env("LANGCHAIN_API_KEY")

    self.client = Client()


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

    self.client.create_run(
      name=name,
      inputs={"input": input},
      outputs={"output": output},
      error=error,
      run_type="chain"
    )