move custom tracer to llm filter (#267)

demos/shared/chatbot_ui/.vscode/launch.json

@@ -15,19 +15,21 @@
         "LLM": "1",
         "CHAT_COMPLETION_ENDPOINT": "http://localhost:10000/v1",
         "STREAMING": "True",
-        "ARCH_CONFIG": "../demos/weather_forecast/arch_config.yaml"
+        "ARCH_CONFIG": "../../weather_forecast/arch_config.yaml"
       }
     },
     {
+      "python": "${workspaceFolder}/venv/bin/python",
       "name": "chatbot-ui llm",
-      "cwd": "${workspaceFolder}/app",
       "type": "debugpy",
       "request": "launch",
-      "program": "run.py",
+      "program": "run_stream.py",
       "console": "integratedTerminal",
       "env": {
         "LLM": "1",
-        "CHAT_COMPLETION_ENDPOINT": "http://localhost:12000/v1"
+        "CHAT_COMPLETION_ENDPOINT": "http://localhost:12000/v1",
+        "STREAMING": "True",
+        "ARCH_CONFIG": "../../llm_routing/arch_config.yaml"
       }
     },
   ]

demos/shared/chatbot_ui/common.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 import json
 import logging
 import os
@@ -159,13 +160,44 @@ def get_prompt_targets():
             config = yaml.safe_load(file)
 
             available_tools = []
-            for target in config["prompt_targets"]:
-                if not target.get("default", False):
-                    available_tools.append(
-                        convert_prompt_target_to_openai_format(target)
-                    )
+            if "prompt_targets" in config:
+                for target in config["prompt_targets"]:
+                    if not target.get("default", False):
+                        available_tools.append(
+                            convert_prompt_target_to_openai_format(target)
+                        )
+
+                return {tool["name"]: tool["info"] for tool in available_tools}
+            elif "llm_providers" in config:
+                return config["llm_providers"]
 
-            return {tool["name"]: tool["info"] for tool in available_tools}
     except Exception as e:
         log.info(e)
         return None
+
+
+def get_llm_models():
+    try:
+        with open(os.getenv("ARCH_CONFIG", "arch_config.yaml"), "r") as file:
+            config = yaml.safe_load(file)
+
+            available_models = [""]
+            default_llm = None
+            for llm_providers in config["llm_providers"]:
+                if llm_providers.get("default", False):
+                    default_llm = llm_providers["name"]
+                else:
+                    available_models.append(llm_providers["name"])
+
+            # place default model at the beginning of the list
+            if default_llm:
+                available_models.insert(0, default_llm)
+            return available_models
+    except Exception as e:
+        log.info(e)
+        return []
+
+
+def format_log(message):
+    time_now = datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
+    return f"{time_now} - {message}"

demos/shared/chatbot_ui/run_stream.py

@@ -8,7 +8,7 @@ from typing import List, Optional, Tuple
 from openai import OpenAI
 from dotenv import load_dotenv
 
-from common import get_prompt_targets, process_stream_chunk
+from common import format_log, get_llm_models, get_prompt_targets, process_stream_chunk
 
 load_dotenv()
 
@@ -36,20 +36,28 @@ CSS_STYLE = """
 footer {visibility: hidden}
 """
 
-client = OpenAI(
-    api_key="--",
-    base_url=CHAT_COMPLETION_ENDPOINT,
-)
-
 
 def chat(
     query: Optional[str],
     conversation: Optional[List[Tuple[str, str]]],
     history: List[dict],
+    debug_output: str,
+    model_selector: str,
 ):
     history.append({"role": "user", "content": query})
 
+    if debug_output is None:
+        debug_output = ""
+
     try:
+        headers = {}
+        if model_selector and model_selector != "":
+            headers["x-arch-llm-provider-hint"] = model_selector
+        client = OpenAI(
+            api_key="--",
+            base_url=CHAT_COMPLETION_ENDPOINT,
+            default_headers=headers,
+        )
         response = client.chat.completions.create(
             # we select model from arch_config file
             model="--",
@@ -65,15 +73,20 @@ def chat(
 
     conversation.append((query, ""))
 
+    model_is_set = False
     for chunk in response:
         tokens = process_stream_chunk(chunk, history)
+        if tokens and not model_is_set:
+            model_is_set = True
+            model = history[-1]["model"]
+            debug_output = debug_output + "\n" + format_log(f"model: {model}")
         if tokens:
             conversation[-1] = (
                 conversation[-1][0],
                 conversation[-1][1] + tokens,
             )
 
-            yield "", conversation, history
+            yield "", conversation, history, debug_output, model_selector
 
 
 def main():
@@ -94,8 +107,17 @@ def main():
                             value=get_prompt_targets(),
                             show_indices=False,
                             elem_classes="json-container",
-                            min_height="95vh",
+                            min_height="50vh",
                         )
+                    model_selector_textbox = gr.Dropdown(
+                        get_llm_models(),
+                        label="override model",
+                        elem_classes="dropdown",
+                    )
+                    debug_output = gr.TextArea(
+                        label="debug output",
+                        elem_classes="debug_output",
+                    )
 
             with gr.Column(scale=2):
                 chatbot = gr.Chatbot(
@@ -110,7 +132,9 @@ def main():
                 )
 
             textbox.submit(
-                chat, [textbox, chatbot, history], [textbox, chatbot, history]
+                chat,
+                [textbox, chatbot, history, debug_output, model_selector_textbox],
+                [textbox, chatbot, history, debug_output, model_selector_textbox],
             )
 
     demo.launch(server_name="0.0.0.0", server_port=8080, show_error=True, debug=True)