feat: Langfuse integration, system-wide tracing configuration (LangSmith&Langfuse)

README.md

@@ -60,6 +60,7 @@ The RAI framework aims to:
 - [x] Integration with LangChain to abstract vendors and access convenient AI tools.
 - [x] Tasks in natural language to nav2 goals.
 - [x] NoMaD integration.
+- [x] Tracing.
 - [ ] Grounded SAM 2 integration.
 - [ ] Improved Human-Robot Interaction with voice and text.
 - [ ] SDK for RAI developers.

config.toml

@@ -17,3 +17,13 @@ simple_model = "llama3.1"
 complex_model = "llama3.1:70b"
 embeddings_model = "llama3.1"
 base_url = "http://localhost:11434"
+
+[tracing]
+project = "rai"
+
+[tracing.langfuse]
+use_langfuse = false
+host = "https://cloud.langfuse.com"
+
+[tracing.langsmith]
+use_langsmith = false

docs/developer_guide.md

@@ -94,6 +94,7 @@ agent.invoke({"messages": ["Please pick up an object and scan it."]})
 
 Additional resources:
 
+- [Tracing](tracing.md) configuration for genAI models and agents.
 - [Beta demos](demos.md).
 - [Multimodal Messages](multimodal_messages.md) definition.
 - Available ROS 2 packages: [ros packages](ros_packages.md).

docs/tracing.md

@@ -0,0 +1,68 @@
+# Tracing Configuration
+
+RAI supports tracing capabilities to help monitor and analyze the performance of your AI models, at a minor performance cost. By default, tracing is off. This document outlines how to configure tracing for your RAI project.
+
+## Configuration
+
+Tracing configuration is managed through the `config.toml` file. The relevant parameters for tracing are:
+
+### Project Name
+
+The `project` field under the `[tracing]` section sets the name for your tracing project. This name will be used to identify your project in the tracing tools.
+
+> [!NOTE]  
+> Project name is currently only used by LangSmith. Langfuse will upload traces to the default project.
+
+### Langfuse (open-source)
+
+[Langfuse](https://langfuse.com/) is an open-source observability & analytics platform for LLM applications.
+
+To enable Langfuse tracing:
+
+1. Set `use_langfuse = true` in the `config.toml` file.
+2. Set the `LANGFUSE_PUBLIC_KEY` and `LANGFUSE_SECRET_KEY` environment variables with your Langfuse credentials.
+3. Optionally, you can specify a custom Langfuse host by modifying the `host` field under `[tracing.langfuse]`.
+
+### LangSmith (closed-source, paid)
+
+[LangSmith](https://www.langchain.com/langsmith) is a platform for debugging, testing, and monitoring LangChain applications.
+
+To enable LangSmith tracing:
+
+1. Set `use_langsmith = true` in the `config.toml` file.
+2. Set the `LANGCHAIN_API_KEY` environment variable with your LangSmith API key.
+
+## Usage
+
+To enable tracing in your RAI application, you need to import the get_tracing_callbacks() function and add it to the configuration when invoking your agent or model. Here's how to do it:
+
+1. First, import the get_tracing_callbacks() function:
+
+```python
+from rai.utils.model_initialization import get_tracing_callbacks
+```
+
+2. Then, add it to the configuration when invoking your agent or model:
+
+```python
+response = agent.invoke(
+    input_dict,
+    config={"callbacks": get_tracing_callbacks()}
+)
+```
+
+By adding the get_tracing_callbacks() to the config parameter, you enable tracing for that specific invocation. The get_tracing_callbacks() function returns a list of callback handlers based on your configuration in config.toml.
+
+## Troubleshooting
+
+If you encounter issues with tracing:
+
+1. Ensure all required environment variables are set correctly.
+2. Check whether tracing is on by checking whether `use_langsmith` or `use_langfuse` flag is set to `true` in `config.toml`.
+3. Verify that you have the necessary permissions and valid API keys for the tracing services you're using.
+4. Look for any error messages in your application logs related to tracing initialization.
+
+For more detailed information on using these tracing tools, refer to their respective documentation:
+
+- [LangSmith Documentation](https://docs.langchain.com/docs/langsmith)
+- [Langfuse Documentation](https://langfuse.com/docs)

src/rai/rai/utils/model_initialization.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from dataclasses import dataclass
-from typing import Literal
+from typing import List, Literal
 
 import tomli
+from langchain_core.callbacks.base import BaseCallbackHandler
 
 
 @dataclass
@@ -40,12 +42,31 @@ class OllamaConfig(ModelConfig):
     base_url: str
 
 
+@dataclass
+class LangfuseConfig:
+    use_langfuse: bool
+    host: str
+
+
+@dataclass
+class LangsmithConfig:
+    use_langsmith: bool
+
+
+@dataclass
+class TracingConfig:
+    project: str
+    langfuse: LangfuseConfig
+    langsmith: LangsmithConfig
+
+
 @dataclass
 class RAIConfig:
     vendor: VendorConfig
     aws: AWSConfig
     openai: ModelConfig
     ollama: OllamaConfig
+    tracing: TracingConfig
 
 
 def load_config() -> RAIConfig:
@@ -56,6 +77,11 @@ def load_config() -> RAIConfig:
         aws=AWSConfig(**config_dict["aws"]),
         openai=ModelConfig(**config_dict["openai"]),
         ollama=OllamaConfig(**config_dict["ollama"]),
+        tracing=TracingConfig(
+            project=config_dict["tracing"]["project"],
+            langfuse=LangfuseConfig(**config_dict["tracing"]["langfuse"]),
+            langsmith=LangsmithConfig(**config_dict["tracing"]["langsmith"]),
+        ),
     )
 
 
@@ -108,3 +134,30 @@ def get_embeddings_model():
         )
     else:
         raise ValueError(f"Unknown embeddings vendor: {vendor}")
+
+
+def get_tracing_callbacks() -> List[BaseCallbackHandler]:
+    config = load_config()
+    callbacks: List[BaseCallbackHandler] = []
+    if config.tracing.langfuse.use_langfuse:
+        from langfuse.callback import CallbackHandler  # type: ignore
+
+        public_key = os.getenv("LANGFUSE_PUBLIC_KEY", None)
+        secret_key = os.getenv("LANGFUSE_SECRET_KEY", None)
+        if public_key is None or secret_key is None:
+            raise ValueError("LANGFUSE_PUBLIC_KEY or LANGFUSE_SECRET_KEY is not set")
+
+        callback = CallbackHandler(
+            public_key=public_key,
+            secret_key=secret_key,
+            host=config.tracing.langfuse.host,
+        )
+        callbacks.append(callback)
+
+    if config.tracing.langsmith.use_langsmith:
+        os.environ["LANGCHAIN_TRACING_V2"] = "true"
+        os.environ["LANGCHAIN_PROJECT"] = config.tracing.project
+        api_key = os.getenv("LANGCHAIN_API_KEY", None)
+        if api_key is None:
+            raise ValueError("LANGCHAIN_API_KEY is not set")
+    return callbacks