NIM not working yet
Some checks failed
Installer CI / smoke-test-on-dev (push) Failing after 5s
Installer CI / lint (push) Failing after 9s

This commit is contained in:
Kai Wu 2025-07-29 14:26:58 -07:00
parent 7065b0fb4d
commit 31a15332c4
3 changed files with 218 additions and 9 deletions

View file

@ -5,21 +5,40 @@
# the root directory of this source tree.
import os
import streamlit as st
from llama_stack_client import LlamaStackClient
class LlamaStackApi:
def __init__(self):
# Initialize provider data from environment variables
self.provider_data = {
"fireworks_api_key": os.environ.get("FIREWORKS_API_KEY", ""),
"together_api_key": os.environ.get("TOGETHER_API_KEY", ""),
"sambanova_api_key": os.environ.get("SAMBANOVA_API_KEY", ""),
"openai_api_key": os.environ.get("OPENAI_API_KEY", ""),
"tavily_search_api_key": os.environ.get("TAVILY_SEARCH_API_KEY", ""),
}
# Check if we have any API keys stored in session state
if st.session_state.get("tavily_search_api_key"):
self.provider_data["tavily_search_api_key"] = st.session_state.get("tavily_search_api_key")
# Initialize the client
self.client = LlamaStackClient(
base_url=os.environ.get("LLAMA_STACK_ENDPOINT", "http://localhost:8321"),
provider_data={
"fireworks_api_key": os.environ.get("FIREWORKS_API_KEY", ""),
"together_api_key": os.environ.get("TOGETHER_API_KEY", ""),
"sambanova_api_key": os.environ.get("SAMBANOVA_API_KEY", ""),
"openai_api_key": os.environ.get("OPENAI_API_KEY", ""),
"tavily_search_api_key": os.environ.get("TAVILY_SEARCH_API_KEY", ""),
},
provider_data=self.provider_data,
)
def update_provider_data(self, key, value):
"""Update a specific provider data key and reinitialize the client"""
self.provider_data[key] = value
# Reinitialize the client with updated provider data
self.client = LlamaStackClient(
base_url=os.environ.get("LLAMA_STACK_ENDPOINT", "http://localhost:8321"),
provider_data=self.provider_data,
)
def run_scoring(self, row, scoring_function_ids: list[str], scoring_params: dict | None):

View file

@ -4,6 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
import os
import streamlit as st
from llama_stack.distribution.ui.modules.api import llama_stack_api
@ -11,6 +12,37 @@ from llama_stack.distribution.ui.modules.api import llama_stack_api
def providers():
st.header("🔍 API Providers")
# API Key Management Section
st.subheader("API Key Management")
# Create a form for API key input
with st.form("api_keys_form"):
# Get the current value from session state or environment variable
tavily_key = st.session_state.get("tavily_search_api_key", os.environ.get("TAVILY_SEARCH_API_KEY", ""))
# Input field for Tavily Search API key
tavily_search_api_key = st.text_input(
"Tavily Search API Key",
value=tavily_key,
type="password",
help="Enter your Tavily Search API key. This will be used for search operations."
)
# Submit button
submit_button = st.form_submit_button("Save API Keys")
if submit_button:
# Store the API key in session state
st.session_state["tavily_search_api_key"] = tavily_search_api_key
# Update the client with the new API key
llama_stack_api.update_provider_data("tavily_search_api_key", tavily_search_api_key)
st.success("API keys saved successfully!")
# Display API Providers
st.subheader("Available API Providers")
apis_providers_lst = llama_stack_api.client.providers.list()
api_to_providers = {}
for api_provider in apis_providers_lst: