W&B Prompts: LLMOps, 언어모델 E2E 대시보드

suite of LLMOps tool built for the development of LLM-powered applications

Weight and Biases는 wandb라는 패키지를 제공하고 있다. 본래 이 라이브러리는 MLOps용으로 TensorBoard와 유사하게 metric들이 학습과정에서 어떻게 변화하는지를 표현해주는 대시보드 기능을 제공한다.

W&B Prompts는 LLM에서 있었던 input, output 그리고 파라미터 값들, 결과가 성공적이었는지/실패였는지 등을 포함하여 편리하게 트래킹할 수 있도록 대시보드를 제공한다.

• Trace Timeline: LLM에서의 각 execution 스텝과 상태를 그래프 형태로 표현, 클릭해서 누르고 보면 좀 더 자세하게 parameter값이나 어디서 에러가 발생했는지 등을 살펴볼 수 있다.

• Trace Table: 모든 trace (input + output, chain, error 등)을 holistic view로 제공한다.

이를 모두 구현하는 코드는 길지 않다. 대신 OpenAI API Key와 Wandb의 ID(API Key)를 요구하니 wandb 계정이 없는 사람은 가입절차를 거쳐야한다. 샘플 코드는 아래에 제공되고 있다:

https://colab.research.google.com/github/wandb/examples/blob/master/colabs/prompts/W%26B_Prompts_Quickstart.ipynb

위 Jupyter 파일에서 코드블럭을 따로 나누지 않고 옮겨적으면 아래와 같다:

!pip install wandb --upgrade
!pip install langchain openai -qqq

os.environ["OPENAI_API_KEY"] = API_KEY

from wandb.integration.langchain import WandbTracer
### This line requires you to have wandb API key
WandbTracer.init({"project": "wandb_prompts_quickstart"})

### LangChain
from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.llms import OpenAI

llm = OpenAI(temperature=0)
tools = load_tools(["llm-math"], llm=llm)
agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)

questions = [
    "Find the square root of 5.4.",
    "What is 3 divided by 7.34 raised to the power of pi?",
    "What is the sin of 0.47 radians, divided by the cube root of 27?",
    "what is 1 divided by zero"
]

### Run Agent
for question in questions:
  try:
    answer = agent.run(question)
    print(answer)
  except Exception as e:
    print(e)
    pass

### Stop watching
WandbTracer.finish()
# > wandb: Finished uploading data to W&B at https:{주소가 팝업됨}

GitHub에 wandb_tracer.py를 참고하면 LangChain에서 지원하고 있는 callback을 활용하고 있음을 확인할 수 있다. Docs에 따르면 Models, Chains, 그리고 Agents의 call에 대한 로그를 남긴다.

https://github.com/wandb/wandb/blob/main/wandb/integration/langchain/wandb_tracer.py

참고로 LangChain에서는 나만의 custom Callback Handler 함수(MyCustomCallbackHandler)를 구현할 수 있도록 해준다.

from typing import Any, Dict, List, Optional, Union

from langchain.agents import initialize_agent, load_tools
from langchain.agents import AgentType
from langchain.callbacks.base import CallbackManager, BaseCallbackHandler
from langchain.llms import OpenAI
from langchain.schema import AgentAction, AgentFinish, LLMResult

class MyCustomCallbackHandler(BaseCallbackHandler):
    """Custom CallbackHandler."""

    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Print out the prompts."""
        pass

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Do nothing."""
        pass

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Do nothing."""
        pass

    def on_llm_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Print out that we are entering a chain."""
        class_name = serialized["name"]
        print(f"\n\n\033[1m> Entering new {class_name} chain...\033[0m")

    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """Print out that we finished a chain."""
        print("\n\033[1m> Finished chain.\033[0m")

    def on_chain_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        """Do nothing."""
        pass

    def on_agent_action(
        self, action: AgentAction, color: Optional[str] = None, **kwargs: Any
    ) -> Any:
        """Run on agent action."""
        print(action)

    def on_tool_end(
        self,
        output: str,
        color: Optional[str] = None,
        observation_prefix: Optional[str] = None,
        llm_prefix: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """If not the final action, print out observation."""
        print(output)

    def on_tool_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_text(
        self,
        text: str,
        color: Optional[str] = None,
        end: str = "",
        **kwargs: Optional[str],
    ) -> None:
        """Run when agent ends."""
        print(text)

    def on_agent_finish(
        self, finish: AgentFinish, color: Optional[str] = None, **kwargs: Any
    ) -> None:
        """Run on agent end."""
        print(finish.log)

manager = CallbackManager([MyCustomCallbackHandler()])
llm = OpenAI(temperature=0, callback_manager=manager, verbose=True)
tools = load_tools(["llm-math", "serpapi"], llm=llm, callback_manager=manager)
agent = initialize_agent(
    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, callback_manager=manager
)
agent.run("Who won the US Open men's final in 2019? What is his age raised to the 0.334 power?")

> Entering new AgentExecutor chain...
AgentAction(tool='Search', tool_input="US Open men's final 2019 winner", log=' I need to find out who won the US Open men\'s final in 2019 and then calculate his age raised to the 0.334 power.\nAction: Search\nAction Input: "US Open men\'s final 2019 winner"')
Rafael Nadal defeated Daniil Medvedev in the final, 7–5, 6–3, 5–7, 4–6, 6–4 to win the men's singles tennis title at the 2019 US Open. It was his fourth US ...
AgentAction(tool='Search', tool_input='Rafael Nadal age', log=' I need to find out the age of the winner\nAction: Search\nAction Input: "Rafael Nadal age"')
36 years
AgentAction(tool='Calculator', tool_input='36^0.334', log=' I now need to calculate his age raised to the 0.334 power\nAction: Calculator\nAction Input: 36^0.334')
Answer: 3.3098250249682484

 I now know the final answer
Final Answer: Rafael Nadal, aged 36, won the US Open men's final in 2019 and his age raised to the 0.334 power is 3.3098250249682484.

> Finished chain.

LLM을 call할 때 LangChain을 사용하지 않더라도 Trace를 남길 수 있는데, 이 내용은 생략하였다.

뿐만 아니라 OpenAI Evals도 함께 통합되어 있어, evaluation job에 대한 모니터링도 가능하다.

만약 production중인 서비스에 결합하고 싶다면, Javascript SDK도 제공한다.