Integrate langsmith for better observability

pyproject.toml

@@ -16,7 +16,7 @@ packages = [{include = "vision_agent"}]
 "documentation" = "https://github.com/landing-ai/vision-agent"
 
 [tool.poetry.dependencies] # main dependency group
-python = ">=3.9"
+python = ">=3.9,<4.0"
 numpy = ">=1.21.0,<2.0.0"
 pillow = "10.*"
 requests = "2.*"
@@ -32,6 +32,7 @@ scipy = "1.13.*"
 nbclient = "^0.10.0"
 nbformat = "^5.10.4"
 rich = "^13.7.1"
+langsmith = "^0.1.58"
 
 [tool.poetry.group.dev.dependencies]
 autoflake = "1.*"

vision_agent/agent/vision_agent_v2.py

@@ -7,6 +7,7 @@
 from rich.console import Console
 from rich.syntax import Syntax
 from tabulate import tabulate
+from langsmith import traceable
 
 from vision_agent.agent import Agent
 from vision_agent.agent.vision_agent_v2_prompt import (
@@ -66,6 +67,7 @@ def extract_json(json_str: str) -> Dict[str, Any]:
  return json_dict # type: ignore
 
 
+@traceable(name="planning")
 def write_plan(
  chat: List[Dict[str, str]],
  plan: Optional[List[Dict[str, Any]]],
@@ -214,6 +216,7 @@ def write_and_exec_code(
  return success, code, result, working_memory
 
 
+@traceable(name="plan execution")
 def run_plan(
  user_req: str,
  plan: List[Dict[str, Any]],
@@ -333,6 +336,7 @@ def __call__(
  results = self.chat_with_workflow(input, image, plan)
  return results["code"] # type: ignore
 
+ @traceable
  def chat_with_workflow(
  self,
  chat: List[Dict[str, str]],

vision_agent/llm/llm.py

@@ -2,7 +2,7 @@
 import os
 from abc import ABC, abstractmethod
 from typing import Any, Callable, Dict, List, Mapping, Optional, Union, cast
-
+from langsmith.wrappers import wrap_openai
 from openai import AzureOpenAI, OpenAI
 
 from vision_agent.tools import (
@@ -41,9 +41,9 @@ def __init__(
  **kwargs: Any
  ):
  if not api_key:
- self.client = OpenAI()
+ self.client = wrap_openai(OpenAI())
  else:
- self.client = OpenAI(api_key=api_key)
+ self.client = wrap_openai(OpenAI(api_key=api_key))
 
  self.model_name = model_name
  self.system_prompt = system_prompt
@@ -165,8 +165,10 @@ def __init__(
  if not azure_endpoint:
  raise ValueError("Azure OpenAI endpoint is required.")
 
- self.client = AzureOpenAI(
- api_key=api_key, api_version=api_version, azure_endpoint=azure_endpoint
+ self.client = wrap_openai(
+ AzureOpenAI(
+ api_key=api_key, api_version=api_version, azure_endpoint=azure_endpoint
+ )
  )
  self.model_name = model_name
  self.kwargs = kwargs