Instrument tracing with langfuse (#167)

* add langsmith

* Update requirements and add async support for LLM completion

- Upgrade ZenML to version 0.74.0
- Add Pinecone, nest_asyncio, and asyncio to requirements
- Implement async version of get_completion_from_messages
- Add Langsmith callback for LLM requests
- Improve error handling for async completion

* Refactor Pinecone and import statements in populate_index and llm_utils

- Reorganize import statements in populate_index.py and llm_utils.py
- Remove redundant Pinecone import in populate_index.py
- Improve code formatting and import order
- Minor code cleanup and optimization

* Instrument langsmith tracing

* Add environment-specific tracing tags for Gradio deployment

- Import os module to handle environment variables
- Add APP_ENVIRONMENT variable with default "dev"
- Include tracing tags in predict function with environment context

Author: strickvl

llm-complete-guide/ZENML_VERSION.txt

@@ -1 +1 @@
-0.71.0
+0.74.0

llm-complete-guide/constants.py

@@ -26,7 +26,7 @@
 # ZenML constants
 ZENML_CHATBOT_MODEL = "zenml-docs-qa-chatbot"
 ZENML_CHATBOT_MODEL_NAME = "zenml-docs-qa-chatbot"
-ZENML_CHATBOT_MODEL_VERSION = "0.71.0-dev"
+ZENML_CHATBOT_MODEL_VERSION = "0.74.0-dev"
 
 # Scraping constants
 RATE_LIMIT = 5  # Maximum number of requests per second

llm-complete-guide/deployment_hf.py

@@ -1,14 +1,20 @@
 import logging
+import os
+import time
 
 import gradio as gr
 from constants import SECRET_NAME
+from langfuse import Langfuse
 from utils.llm_utils import process_input_with_retrieval
 from zenml.client import Client
 
-# Set up logging
+langfuse = Langfuse()
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+APP_ENVIRONMENT = os.getenv("GRADIO_ZENML_APP_ENVIRONMENT", "dev")
+
 # Initialize ZenML client and verify secret access
 try:
     client = Client()
@@ -21,6 +27,80 @@
     raise RuntimeError(f"Application startup failed: {e}")
 
 
+def get_langfuse_trace_id() -> str | None:
+    """Get the trace from Langfuse.
+
+    This is a very naive implementation. It simply returns the id of the first trace
+    in the last 60 seconds. Will retry up to 3 times if no traces are found or if
+    there's an error.
+
+    Returns:
+        str | None: The trace ID if found, None otherwise
+    """
+    logger.info("Getting trace from Langfuse")
+    retries = 0
+    max_retries = 3
+    while retries < max_retries:
+        try:
+            # Wait 5 seconds before making the API call
+            time.sleep(5)
+            traces = langfuse.fetch_traces(
+                limit=1, order_by="timestamp.desc"
+            ).data
+            if not traces:
+                retries += 1
+                if retries == max_retries:
+                    logger.error(
+                        f"No traces found after {max_retries} attempts"
+                    )
+                    return None
+                logger.warning(
+                    f"No traces found (attempt {retries}/{max_retries})"
+                )
+                time.sleep(10)
+                continue
+            return traces[0].id
+        except Exception as e:
+            retries += 1
+            if retries == max_retries:
+                logger.error(
+                    f"Error fetching traces after {max_retries} attempts: {e}"
+                )
+                return None
+            logger.warning(
+                f"Error fetching traces (attempt {retries}/{max_retries}): {e}"
+            )
+            time.sleep(10)
+    return None
+
+
+def vote(data: gr.LikeData):
+    """Vote on a response.
+
+    Args:
+        data (gr.LikeData): The vote data.
+    """
+
+    trace_id = get_langfuse_trace_id()
+    logger.info(f"Vote data: {data}")
+    if data.liked:
+        logger.info("Vote up")
+        langfuse.score(
+            trace_id=trace_id,
+            name="user-explicit-feedback",
+            value="like",
+            comment="I like this response",
+        )
+    else:
+        logger.info("Vote down")
+        langfuse.score(
+            trace_id=trace_id,
+            name="user-explicit-feedback",
+            value="dislike",
+            comment="I don't like the response",
+        )
+
+
 def predict(message, history):
     try:
         # add the prod flag here
@@ -29,18 +109,29 @@ def predict(message, history):
             n_items_retrieved=20,
             use_reranking=True,
             model_version_stage="production",
+            tracing_tags=["gradio", "web-interface", APP_ENVIRONMENT],
         )
     except Exception as e:
         logger.error(f"Error processing message: {e}")
         return f"Sorry, I encountered an error: {str(e)}"
 
 
-# Launch the Gradio interface
-interface = gr.ChatInterface(
-    predict,
-    title="ZenML Documentation Assistant",
-    description="Ask me anything about ZenML!",
-)
+with gr.Blocks() as interface:
+    custom_chatbot = gr.Chatbot(
+        type="messages",
+        editable=True,
+    )
+
+    gr.ChatInterface(
+        predict,
+        type="messages",
+        title="ZenML Documentation Assistant",
+        description="Ask me anything about ZenML!",
+        chatbot=custom_chatbot,
+        theme="shivi/calm_seafoam",
+    )
+
+    custom_chatbot.like(vote, None, None)
 
 if __name__ == "__main__":
     interface.launch(server_name="0.0.0.0", share=False)

llm-complete-guide/requirements.txt

@@ -1,4 +1,4 @@
-zenml[server]>=0.73.0
+git+https://github.com/zenml-io/zenml.git@develop#egg=zenml[server] # will work for zenml>=0.75.0
 ratelimit
 pgvector
 psycopg2-binary
@@ -18,10 +18,12 @@ pyarrow
 rerankers[flashrank]
 datasets
 torch
-gradio
+gradio>=5.13.0
 huggingface-hub
 elasticsearch
 tenacity
+langfuse
+pinecone
 
 # optional requirements for S3 artifact store
 # s3fs>2022.3.0

llm-complete-guide/run.py

@@ -234,7 +234,11 @@ def main(
             )
         # add the prod flag here
         response = process_input_with_retrieval(
-            query_text, model=model, use_reranking=use_reranker, model_version_stage="production"
+            query_text,
+            model=model,
+            use_reranking=use_reranker,
+            model_version_stage="production",
+            tracing_tags=["cli", "dev"],
         )
         console = Console()
         md = Markdown(response)

llm-complete-guide/steps/populate_index.py

@@ -48,12 +48,16 @@
 from PIL import Image, ImageDraw, ImageFont
 from sentence_transformers import SentenceTransformer
 from structures import Document
-from utils.llm_utils import get_db_conn, get_es_client, get_pinecone_client, split_documents
+from utils.llm_utils import (
+    get_db_conn,
+    get_es_client,
+    get_pinecone_client,
+    split_documents,
+)
 from zenml import ArtifactConfig, get_step_context, log_metadata, step
 from zenml.client import Client
 from zenml.metadata.metadata_types import Uri
-import pinecone
-from pinecone import Pinecone, ServerlessSpec
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -642,7 +646,7 @@ def index_generator(
         documents (str): JSON string containing the documents to index.
         index_type (IndexType, optional): Type of index to generate. Defaults to IndexType.POSTGRES.
     """
-    # get model version 
+    # get model version
     context = get_step_context()
     model_version_stage = context.model_version.stage
     if index_type == IndexType.ELASTICSEARCH:
@@ -825,7 +829,9 @@ def _index_generator_postgres(documents: str) -> None:
             conn.close()
 
 
-def _index_generator_pinecone(documents: str, model_version_stage: str) -> None:
+def _index_generator_pinecone(
+    documents: str, model_version_stage: str
+) -> None:
     """Generates a Pinecone index for the given documents.
 
     Args:
@@ -856,8 +862,8 @@ def _index_generator_pinecone(documents: str, model_version_stage: str) -> None:
                 "parent_section": doc["parent_section"] or "",
                 "url": doc["url"],
                 "page_content": doc["page_content"],
-                "token_count": doc["token_count"]
-            }
+                "token_count": doc["token_count"],
+            },
         }
         batch.append(vector_record)
 
@@ -870,7 +876,9 @@ def _index_generator_pinecone(documents: str, model_version_stage: str) -> None:
     if batch:
         index.upsert(vectors=batch)
 
-    logger.info(f"Successfully indexed {len(docs)} documents to Pinecone index")
+    logger.info(
+        f"Successfully indexed {len(docs)} documents to Pinecone index"
+    )
 
 
 def _log_metadata(index_type: IndexType) -> None:
@@ -914,7 +922,9 @@ def _log_metadata(index_type: IndexType) -> None:
         store_name = "pinecone"
         connection_details = {
             "api_key": "**********",
-            "environment": client.get_secret(SECRET_NAME_PINECONE).secret_values["pinecone_env"],
+            "environment": client.get_secret(
+                SECRET_NAME_PINECONE
+            ).secret_values["pinecone_env"],
         }
 
     log_metadata(