choose_mode_instrumented.py

choose_mode_instrumented.py

@action(reads=["prompt"], writes=["mode"])
def choose_mode(state: State, __tracer) -> Tuple[dict, State]:
    prompt = ...
    with __tracer("query_openai") as tracer:
      result = _get_openai_client().chat.completions.create(
          model="gpt-4",
          messages=[
              {"role": "system", "content": "You are a helpful assistant"},
              {"role": "user", "content": prompt},
          ],
      )
      tracer.log_attributes(prompt=prompt, model="gpt-4", result=json.dumps(result))
    ... # continue as normal