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):