.

Author: rguan72

backend/onyx/evals/eval_cli.py

@@ -5,7 +5,6 @@
 
 import argparse
 import json
-import logging
 import os
 from typing import Any
 
@@ -121,7 +120,6 @@ def run_remote(
 
 
 def main() -> None:
-    logging.basicConfig(level=logging.INFO)
     """Main CLI entry point."""
     parser = argparse.ArgumentParser(
         description="Run evaluations with local configurations"

backend/onyx/llm/chat_llm.py

@@ -26,6 +26,7 @@
 from langchain_core.prompt_values import PromptValue
 from litellm.utils import get_supported_openai_params
 
+from onyx.configs.app_configs import BRAINTRUST_ENABLED
 from onyx.configs.app_configs import LOG_DANSWER_MODEL_INTERACTIONS
 from onyx.configs.app_configs import MOCK_LLM_RESPONSE
 from onyx.configs.chat_configs import QA_TIMEOUT
@@ -49,9 +50,8 @@
 litellm.drop_params = True
 litellm.telemetry = False
 
-# if BRAINTRUST_ENABLED:
-#     print("Setting LiteLLM callbacks to braintrust")
-#     litellm.callbacks = ["braintrust"]
+if BRAINTRUST_ENABLED:
+    litellm.callbacks = ["braintrust"]
 
 _LLM_PROMPT_LONG_TERM_LOG_CATEGORY = "llm_prompt"
 VERTEX_CREDENTIALS_FILE_KWARG = "vertex_credentials"

backend/onyx/llm/interfaces.py

@@ -4,6 +4,7 @@
 
 from braintrust import traced
 from langchain.schema.language_model import LanguageModelInput
+from langchain_core.messages import AIMessageChunk
 from langchain_core.messages import BaseMessage
 from pydantic import BaseModel
 
@@ -33,32 +34,29 @@ class LLMConfig(BaseModel):
 
 
 def log_prompt(prompt: LanguageModelInput) -> None:
-    # if isinstance(prompt, list):
-    #     for ind, msg in enumerate(prompt):
-    #         if isinstance(msg, AIMessageChunk):
-    #             if msg.content:
-    #                 log_msg = msg.content
-    #             elif msg.tool_call_chunks:
-    #                 # log_msg = "Tool Calls: " + str(
-    #                 #     [
-    #                 #         {
-    #                 #             key: value
-    #                 #             for key, value in tool_call.items()
-    #                 #             if key != "index"
-    #                 #         }
-    #                 #         for tool_call in msg.tool_call_chunks
-    #                 #     ]
-    #                 # )
-    #                 pass
-    #             else:
-    #                 pass
-    #             # logger.debug(f"Message {ind}:\n{log_msg}")
-    #         else:
-    #             # logger.debug(f"Message {ind}:\n{msg.content}")
-    #             pass
+    if isinstance(prompt, list):
+        for ind, msg in enumerate(prompt):
+            if isinstance(msg, AIMessageChunk):
+                if msg.content:
+                    log_msg = msg.content
+                elif msg.tool_call_chunks:
+                    log_msg = "Tool Calls: " + str(
+                        [
+                            {
+                                key: value
+                                for key, value in tool_call.items()
+                                if key != "index"
+                            }
+                            for tool_call in msg.tool_call_chunks
+                        ]
+                    )
+                else:
+                    pass
+                logger.debug(f"Message {ind}:\n{log_msg}")
+            else:
+                logger.debug(f"Message {ind}:\n{msg.content}")
     if isinstance(prompt, str):
-        # logger.debug(f"Prompt:\n{prompt}")
-        pass
+        logger.debug(f"Prompt:\n{prompt}")
 
 
 class LLM(abc.ABC):

backend/onyx/main.py

@@ -35,6 +35,7 @@
 from onyx.configs.app_configs import APP_PORT
 from onyx.configs.app_configs import AUTH_RATE_LIMITING_ENABLED
 from onyx.configs.app_configs import AUTH_TYPE
+from onyx.configs.app_configs import BRAINTRUST_ENABLED
 from onyx.configs.app_configs import DISABLE_GENERATIVE_AI
 from onyx.configs.app_configs import LOG_ENDPOINT_LATENCY
 from onyx.configs.app_configs import OAUTH_CLIENT_ID
@@ -51,6 +52,7 @@
 from onyx.db.engine.connection_warmup import warm_up_connections
 from onyx.db.engine.sql_engine import get_session_with_current_tenant
 from onyx.db.engine.sql_engine import SqlEngine
+from onyx.evals.tracing import setup_braintrust
 from onyx.file_store.file_store import get_default_file_store
 from onyx.server.api_key.api import router as api_key_router
 from onyx.server.auth_check import check_router_auth
@@ -252,6 +254,10 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     if DISABLE_GENERATIVE_AI:
         logger.notice("Generative AI Q&A disabled")
 
+    if BRAINTRUST_ENABLED:
+        setup_braintrust()
+        logger.notice("Braintrust tracing initialized")
+
     # fill up Postgres connection pools
     await warm_up_connections()