Use Run object from SDK

langchain/callbacks/tracers/langchain.py

@@ -11,7 +11,11 @@
 from langchainplus_sdk import LangChainPlusClient
 
 from langchain.callbacks.tracers.base import BaseTracer
-from langchain.callbacks.tracers.schemas import Run, RunTypeEnum, TracerSession
+from langchain.callbacks.tracers.schemas import (
+    Run,
+    RunTypeEnum,
+    TracerSession,
+)
 from langchain.env import get_runtime_environment
 from langchain.schema import BaseMessage, messages_to_dict
 

langchain/callbacks/tracers/schemas.py

@@ -2,13 +2,13 @@
 from __future__ import annotations
 
 import datetime
-from enum import Enum
 from typing import Any, Dict, List, Optional
 from uuid import UUID
 
+from langchainplus_sdk.schemas import RunBase as BaseRunV2
+from langchainplus_sdk.schemas import RunTypeEnum
 from pydantic import BaseModel, Field, root_validator
 
-from langchain.env import get_runtime_environment
 from langchain.schema import LLMResult
 
 
@@ -88,36 +88,11 @@ class ToolRun(BaseRun):
 # Begin V2 API Schemas
 
 
-class RunTypeEnum(str, Enum):
-    """Enum for run types."""
+class Run(BaseRunV2):
+    """Run schema for the V2 API in the Tracer."""
 
-    tool = "tool"
-    chain = "chain"
-    llm = "llm"
-
-
-class RunBase(BaseModel):
-    """Base Run schema."""
-
-    id: Optional[UUID]
-    start_time: datetime.datetime = Field(default_factory=datetime.datetime.utcnow)
-    end_time: datetime.datetime = Field(default_factory=datetime.datetime.utcnow)
-    extra: Optional[Dict[str, Any]] = None
-    error: Optional[str]
     execution_order: int
-    child_execution_order: Optional[int]
-    serialized: dict
-    inputs: dict
-    outputs: Optional[dict]
-    reference_example_id: Optional[UUID]
-    run_type: RunTypeEnum
-    parent_run_id: Optional[UUID]
-
-
-class Run(RunBase):
-    """Run schema when loading from the DB."""
-
-    name: str
+    child_execution_order: int
     child_runs: List[Run] = Field(default_factory=list)
 
     @root_validator(pre=True)
@@ -131,26 +106,19 @@ def assign_name(cls, values: dict) -> dict:
         return values
 
 
-class RunCreate(RunBase):
-    name: str
-    session_name: Optional[str] = None
-
-    @root_validator(pre=True)
-    def add_runtime_env(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        """Add env info to the run."""
-        extra = values.get("extra", {})
-        extra["runtime"] = get_runtime_environment()
-        values["extra"] = extra
-        return values
-
-
-class RunUpdate(BaseModel):
-    end_time: Optional[datetime.datetime]
-    error: Optional[str]
-    outputs: Optional[dict]
-    parent_run_id: Optional[UUID]
-    reference_example_id: Optional[UUID]
-
-
 ChainRun.update_forward_refs()
 ToolRun.update_forward_refs()
+
+__all__ = [
+    "BaseRun",
+    "ChainRun",
+    "LLMRun",
+    "Run",
+    "RunTypeEnum",
+    "ToolRun",
+    "TracerSession",
+    "TracerSessionBase",
+    "TracerSessionV1",
+    "TracerSessionV1Base",
+    "TracerSessionV1Create",
+]

pyproject.toml

@@ -105,7 +105,7 @@ singlestoredb = {version = "^0.6.1", optional = true}
 pyspark = {version = "^3.4.0", optional = true}
 tigrisdb = {version = "^1.0.0b6", optional = true}
 nebula3-python = {version = "^3.4.0", optional = true}
-langchainplus-sdk = ">=0.0.7"
+langchainplus-sdk = ">=0.0.9"
 awadb = {version = "^0.3.2", optional = true}
 azure-search-documents = {version = "11.4.0a20230509004", source = "azure-sdk-dev", optional = true}
 

tests/unit_tests/callbacks/test_callback_manager.py

@@ -203,8 +203,10 @@ def test_callback_manager_inheritance() -> None:
     assert child_manager2.inheritable_handlers == [handler1]
 
 
-def test_callback_manager_configure() -> None:
+def test_callback_manager_configure(monkeypatch: pytest.MonkeyPatch) -> None:
     """Test callback manager configuration."""
+    monkeypatch.setenv("LANGCHAIN_TRACING_V2", "false")
+    monkeypatch.setenv("LANGCHAIN_TRACING", "false")
     handler1, handler2, handler3, handler4 = (
         FakeCallbackHandler(),
         FakeCallbackHandler(),

tests/unit_tests/callbacks/test_schemas.py

@@ -0,0 +1,27 @@
+import langchain.callbacks.tracers.schemas as schemas
+from langchain.callbacks.tracers.schemas import __all__ as schemas_all
+
+
+def test_public_api() -> None:
+    """Test for changes in the public API."""
+    expected_all = [
+        "BaseRun",
+        "ChainRun",
+        "LLMRun",
+        "Run",
+        "RunTypeEnum",
+        "ToolRun",
+        "TracerSession",
+        "TracerSessionBase",
+        "TracerSessionV1",
+        "TracerSessionV1Base",
+        "TracerSessionV1Create",
+    ]
+
+    assert sorted(schemas_all) == expected_all
+
+    # Assert that the object is actually present in the schema module
+    for module_name in expected_all:
+        assert (
+            hasattr(schemas, module_name) and getattr(schemas, module_name) is not None
+        )