add tracing support to langchain

langchain/__init__.py

@@ -4,7 +4,11 @@
 
 from langchain.agents import MRKLChain, ReActChain, SelfAskWithSearchChain
 from langchain.cache import BaseCache
-from langchain.callbacks import set_default_callback_manager, set_handler
+from langchain.callbacks import (
+    set_default_callback_manager,
+    set_handler,
+    set_tracing_callback_manager,
+)
 from langchain.chains import (
     ConversationChain,
     LLMBashChain,
@@ -68,4 +72,5 @@
     "QAWithSourcesChain",
     "PALChain",
     "set_handler",
+    "set_tracing_callback_manager",
 ]

langchain/agents/agent.py

@@ -284,7 +284,7 @@ def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]:
                     observation = tool.func(output.tool_input)
                     color = color_mapping[output.tool]
                     return_direct = tool.return_direct
-                except Exception as e:
+                except (KeyboardInterrupt, Exception) as e:
                     self.callback_manager.on_tool_error(e, verbose=self.verbose)
                     raise e
             else:

langchain/callbacks/__init__.py

@@ -1,11 +1,13 @@
 """Callback handlers that allow listening to events in LangChain."""
+import os
 from contextlib import contextmanager
-from typing import Generator
+from typing import Generator, Optional
 
 from langchain.callbacks.base import BaseCallbackHandler, BaseCallbackManager
 from langchain.callbacks.openai_info import OpenAICallbackHandler
 from langchain.callbacks.shared import SharedCallbackManager
 from langchain.callbacks.stdout import StdOutCallbackHandler
+from langchain.callbacks.tracers import SharedLangChainTracer
 
 
 def get_callback_manager() -> BaseCallbackManager:
@@ -21,7 +23,31 @@ def set_handler(handler: BaseCallbackHandler) -> None:
 
 def set_default_callback_manager() -> None:
     """Set default callback manager."""
-    set_handler(StdOutCallbackHandler())
+    default_handler = os.environ.get("LANGCHAIN_HANDLER", "stdout")
+    if default_handler == "stdout":
+        set_handler(StdOutCallbackHandler())
+    elif default_handler == "langchain":
+        session = os.environ.get("LANGCHAIN_SESSION")
+        set_tracing_callback_manager(session)
+    else:
+        raise ValueError(
+            f"LANGCHAIN_HANDLER should be one of `stdout` "
+            f"or `langchain`, got {default_handler}"
+        )
+
+
+def set_tracing_callback_manager(session_name: Optional[str] = None) -> None:
+    """Set tracing callback manager."""
+    handler = SharedLangChainTracer()
+    callback = get_callback_manager()
+    callback.set_handlers([handler, StdOutCallbackHandler()])
+    if session_name is None:
+        handler.load_default_session()
+    else:
+        try:
+            handler.load_session(session_name)
+        except Exception:
+            raise ValueError(f"session {session_name} not found")
 
 
 @contextmanager

langchain/callbacks/base.py

@@ -1,25 +1,34 @@
 """Base callback handler that can be used to handle callbacks from langchain."""
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List
-
-from pydantic import BaseModel
+from typing import Any, Dict, List, Union
 
 from langchain.schema import AgentAction, AgentFinish, LLMResult
 
 
-class BaseCallbackHandler(BaseModel, ABC):
+class BaseCallbackHandler(ABC):
     """Base callback handler that can be used to handle callbacks from langchain."""
 
-    ignore_llm: bool = False
-    ignore_chain: bool = False
-    ignore_agent: bool = False
-
     @property
     def always_verbose(self) -> bool:
         """Whether to call verbose callbacks even if verbose is False."""
         return False
 
+    @property
+    def ignore_llm(self) -> bool:
+        """Whether to ignore LLM callbacks."""
+        return False
+
+    @property
+    def ignore_chain(self) -> bool:
+        """Whether to ignore chain callbacks."""
+        return False
+
+    @property
+    def ignore_agent(self) -> bool:
+        """Whether to ignore agent callbacks."""
+        return False
+
     @abstractmethod
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
@@ -31,7 +40,9 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Run when LLM ends running."""
 
     @abstractmethod
-    def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Run when LLM errors."""
 
     @abstractmethod
@@ -45,7 +56,9 @@ def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         """Run when chain ends running."""
 
     @abstractmethod
-    def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_chain_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Run when chain errors."""
 
     @abstractmethod
@@ -59,7 +72,9 @@ def on_tool_end(self, output: str, **kwargs: Any) -> None:
         """Run when tool ends running."""
 
     @abstractmethod
-    def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_tool_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Run when tool errors."""
 
     @abstractmethod
@@ -82,15 +97,21 @@ def add_handler(self, callback: BaseCallbackHandler) -> None:
     def remove_handler(self, handler: BaseCallbackHandler) -> None:
         """Remove a handler from the callback manager."""
 
-    @abstractmethod
     def set_handler(self, handler: BaseCallbackHandler) -> None:
         """Set handler as the only handler on the callback manager."""
+        self.set_handlers([handler])
+
+    @abstractmethod
+    def set_handlers(self, handlers: List[BaseCallbackHandler]) -> None:
+        """Set handlers as the only handlers on the callback manager."""
 
 
 class CallbackManager(BaseCallbackManager):
     """Callback manager that can be used to handle callbacks from langchain."""
 
-    handlers: List[BaseCallbackHandler]
+    def __init__(self, handlers: List[BaseCallbackHandler]) -> None:
+        """Initialize callback manager."""
+        self.handlers: List[BaseCallbackHandler] = handlers
 
     def on_llm_start(
         self,
@@ -115,7 +136,10 @@ def on_llm_end(
                     handler.on_llm_end(response)
 
     def on_llm_error(
-        self, error: Exception, verbose: bool = False, **kwargs: Any
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        verbose: bool = False,
+        **kwargs: Any
     ) -> None:
         """Run when LLM errors."""
         for handler in self.handlers:
@@ -146,7 +170,10 @@ def on_chain_end(
                     handler.on_chain_end(outputs)
 
     def on_chain_error(
-        self, error: Exception, verbose: bool = False, **kwargs: Any
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        verbose: bool = False,
+        **kwargs: Any
     ) -> None:
         """Run when chain errors."""
         for handler in self.handlers:
@@ -175,7 +202,10 @@ def on_tool_end(self, output: str, verbose: bool = False, **kwargs: Any) -> None
                     handler.on_tool_end(output, **kwargs)
 
     def on_tool_error(
-        self, error: Exception, verbose: bool = False, **kwargs: Any
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        verbose: bool = False,
+        **kwargs: Any
     ) -> None:
         """Run when tool errors."""
         for handler in self.handlers:
@@ -206,6 +236,6 @@ def remove_handler(self, handler: BaseCallbackHandler) -> None:
         """Remove a handler from the callback manager."""
         self.handlers.remove(handler)
 
-    def set_handler(self, handler: BaseCallbackHandler) -> None:
-        """Set handler as the only handler on the callback manager."""
-        self.handlers = [handler]
+    def set_handlers(self, handlers: List[BaseCallbackHandler]) -> None:
+        """Set handlers as the only handlers on the callback manager."""
+        self.handlers = handlers

langchain/callbacks/openai_info.py

@@ -1,5 +1,5 @@
 """Callback Handler that prints to std out."""
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import AgentAction, AgentFinish, LLMResult
@@ -29,7 +29,9 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
                 if "total_tokens" in token_usage:
                     self.total_tokens += token_usage["total_tokens"]
 
-    def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass
 
@@ -43,7 +45,9 @@ def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         """Print out that we finished a chain."""
         pass
 
-    def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_chain_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass
 
@@ -68,7 +72,9 @@ def on_tool_end(
         """If not the final action, print out observation."""
         pass
 
-    def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_tool_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass
 

langchain/callbacks/shared.py

@@ -1,7 +1,7 @@
 """A shared CallbackManager."""
 
 import threading
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Union
 
 from langchain.callbacks.base import (
     BaseCallbackHandler,
@@ -46,7 +46,9 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         with self._lock:
             self._callback_manager.on_llm_end(response, **kwargs)
 
-    def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Run when LLM errors."""
         with self._lock:
             self._callback_manager.on_llm_error(error, **kwargs)
@@ -63,7 +65,9 @@ def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         with self._lock:
             self._callback_manager.on_chain_end(outputs, **kwargs)
 
-    def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_chain_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Run when chain errors."""
         with self._lock:
             self._callback_manager.on_chain_error(error, **kwargs)
@@ -80,7 +84,9 @@ def on_tool_end(self, output: str, **kwargs: Any) -> None:
         with self._lock:
             self._callback_manager.on_tool_end(output, **kwargs)
 
-    def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_tool_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Run when tool errors."""
         with self._lock:
             self._callback_manager.on_tool_error(error, **kwargs)
@@ -105,7 +111,7 @@ def remove_handler(self, callback: BaseCallbackHandler) -> None:
         with self._lock:
             self._callback_manager.remove_handler(callback)
 
-    def set_handler(self, handler: BaseCallbackHandler) -> None:
-        """Set handler as the only handler on the callback manager."""
+    def set_handlers(self, handlers: List[BaseCallbackHandler]) -> None:
+        """Set handlers as the only handlers on the callback manager."""
         with self._lock:
-            self._callback_manager.handlers = [handler]
+            self._callback_manager.handlers = handlers

langchain/callbacks/stdout.py

@@ -1,5 +1,5 @@
 """Callback Handler that prints to std out."""
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.input import print_text
@@ -19,7 +19,9 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Do nothing."""
         pass
 
-    def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass
 
@@ -34,7 +36,9 @@ 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: Exception, **kwargs: Any) -> None:
+    def on_chain_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass
 
@@ -61,7 +65,9 @@ def on_tool_end(
         print_text(output, color=color)
         print_text(f"\n{llm_prefix}")
 
-    def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_tool_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass
 

langchain/callbacks/streamlit.py

@@ -1,5 +1,5 @@
 """Callback Handler that logs to streamlit."""
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import streamlit as st
 
@@ -22,7 +22,9 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Do nothing."""
         pass
 
-    def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass
 
@@ -37,7 +39,9 @@ def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         """Print out that we finished a chain."""
         st.write("Finished chain.")
 
-    def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_chain_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass
 
@@ -62,7 +66,9 @@ def on_tool_end(
         st.write(f"{observation_prefix}{output}")
         st.write(llm_prefix)
 
-    def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
+    def on_tool_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
         """Do nothing."""
         pass