Merge pull request #8 from dristysrivastava/dristy-safe-infer-app-update

Updated Streamlit app based on updated API, also removed models selection from sidebar

Author: gr8nishan

safe_infer_chatbot_app/README.md

@@ -51,6 +51,9 @@ Then open your browser to `http://localhost:8501`
    ```bash
    export PEBBLO_API_KEY="pebblo-api-key"
    export PROXIMA_HOST="http://your-proxima-host"
+   export MODEL="model-name"
+   export X_PEBBLO_USER="user-email"
+   export MODEL_NAME="model-display-name"
    ```
 
 4. **Run the application**:
@@ -66,31 +69,26 @@ Then open your browser to `http://localhost:8501`
 
 - `PROXIMA_HOST`: Base URL for the SafeInfer API (default: `http://localhost`)
 - `PEBBLO_API_KEY`: Pebblo API Key
+-  `MODEL`: Model Name
+- `MODEL_NAME`: Model Display Name
+- `X_PEBBLO_USER`: User Email
 
 ### API Configuration
 
 The application automatically configures the following endpoints:
 - **Responses**: `{PROXIMA_HOST}/safe_infer/llm/v1/responses`
 - **Health Check**: `{PROXIMA_HOST}/safe_infer/healthz`
 
-### Available Models
-
-- `gpt-4o-mini`: Faster, more cost-effective model
-- `gpt-4o`: Full GPT-4o model with enhanced capabilities
-
 ## 📖 Usage Guide
 
 ### Starting a Conversation
 
-1. **Select a Model**: Choose your preferred model from the sidebar
-2. **Enter API Key** (if required): Add your API key in the configuration section
-3. **Test Connection**: Use the "Test API Connection" button to verify connectivity
-4. **Start Chatting**: Type your message and click "Send" or press Enter
+1. **Test Connection**: Use the "Test API Connection" button to verify connectivity
+2. **Start Chatting**: Type your message and click "Send" or press Enter
 
 ### Chat Features
 
 - **Send Messages**: Type in the text area and click "🚀 Send"
-- **Regenerate Responses**: Click "🔄 Regenerate Last Response" to get a new AI response
 - **Clear History**: Use "Clear Chat History" to start fresh
 - **Export Chat**: Download your conversation as a JSON file
 

safe_infer_chatbot_app/safe_infer_chatbot.py

@@ -4,12 +4,13 @@
 import json
 from typing import Dict, Any
 import time
+from openai import OpenAI
 
 
 
 # Page configuration
 st.set_page_config(
-    page_title="SafeInfer LLM Chatbot",
+    page_title="Finance Ops Chatbot",
     page_icon="🛡️",
     layout="wide",
     initial_sidebar_state="expanded"
@@ -59,22 +60,27 @@
 </style>
 """, unsafe_allow_html=True)
 
-from utils import get_available_models
 
 # API Configuration
 API_KEY = os.getenv("PEBBLO_API_KEY", "")
 API_BASE_URL = os.getenv("PROXIMA_HOST", "http://localhost")
-RESPONSE_API_ENDPOINT = f"{API_BASE_URL}/safe_infer/llm/v1/responses"
+USER_EMAIL = os.getenv("USER_EMAIL", "User")
+USER_TEAM = os.getenv("USER_TEAM", "Finance Ops")
+RESPONSE_API_ENDPOINT = f"{API_BASE_URL}/safe_infer/llm/v1/"
 LLM_PROVIDER_API_ENDPOINT = f"{API_BASE_URL}/api/llm/provider"
-AVAILABLE_MODELS, DEFAULT_MODEL = get_available_models()
+SELECTED_MODEL = os.getenv("MODEL")
+X_PEBBLO_USER = os.getenv("X_PEBBLO_USER", None)
+MODEL_NAME = os.getenv("MODEL_NAME", SELECTED_MODEL)
 
 # Initialize session state
 if 'chat_history' not in st.session_state:
     st.session_state.chat_history = []
 if 'selected_model' not in st.session_state:
-    st.session_state.selected_model = DEFAULT_MODEL
+    st.session_state.selected_model = SELECTED_MODEL
 if 'api_key' not in st.session_state:
     st.session_state.api_key = API_KEY
+if 'model_name' not in st.session_state:
+    st.session_state.model_name = MODEL_NAME
 
 def test_api_connection() -> Dict[str, Any]:
     """Test the API connection"""
@@ -89,76 +95,23 @@ def test_api_connection() -> Dict[str, Any]:
     except Exception as e:
         return {"status": "error", "message": f"Error: {str(e)}"}
 
-def call_safe_infer_api(message: str, model: str, api_key: str = "") -> Dict[str, Any]:
-    """Call the SafeInfer API"""
-    headers = {
-        "Content-Type": "application/json"
-    }
-    
-    if api_key:
-        headers["Authorization"] = f"Bearer {api_key}"
-    
-    payload = {
-        "model": model,
-        "input": message
-    }
-    
+def call_open_ai(message: str, model: str, api_key: str = "") -> Dict[str, Any]:
     try:
-        response = requests.post(
-            RESPONSE_API_ENDPOINT,
-            json=payload,
-            headers=headers,
-            timeout=30
+        default_headers = {"X-PEBBLO-USER": X_PEBBLO_USER} if X_PEBBLO_USER else None
+        client = OpenAI(
+            base_url=RESPONSE_API_ENDPOINT,
+            api_key=api_key,
+            default_headers=default_headers
+        )
+        response = client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": message}]
         )
 
-        if response.status_code == 200:
-            return {"status": "success", "data": response.json()}
-        else:
-            return {
-                "status": "error", 
-                "message": f"API Error {response.status_code}: {response.text}"
-            }
-    except requests.exceptions.Timeout:
-        return {"status": "error", "message": "Request timed out"}
-    except requests.exceptions.ConnectionError:
-        return {"status": "error", "message": "Cannot connect to API"}
+        return {"status": "success", "data": response.choices[0].message.content}
     except Exception as e:
         return {"status": "error", "message": f"Error: {str(e)}"}
 
-def extract_response_content(api_response: Dict[str, Any]) -> str:
-    """Extract the response content from the API response"""
-    try:
-        # Handle different response formats
-        if 'response' in api_response:
-            response_data = api_response['response']
-            if isinstance(response_data, dict):
-                if 'message' in response_data:
-                    if isinstance(response_data['message'], str):
-                        return response_data['message']
-                    elif isinstance(response_data['message'], dict) and 'content' in response_data['message']:
-                        return response_data['message']['content']
-                elif 'content' in response_data:
-                    return response_data['content']
-            elif isinstance(response_data, str):
-                return response_data
-        
-        # Check for direct content
-        elif 'content' in api_response:
-            return api_response['content']
-        
-        # Check for message field
-        elif 'message' in api_response:
-            if isinstance(api_response['message'], str):
-                return api_response['message']
-            elif isinstance(api_response['message'], dict) and 'content' in api_response['message']:
-                return api_response['message']['content']
-        
-        # If none of the above, return the full response as JSON
-        return json.dumps(api_response, indent=2)
-        
-    except Exception as e:
-        return f"Error parsing response: {str(e)}"
-
 def display_chat_message(role: str, content: str, model: str = "", timestamp: str = ""):
     """Display a chat message with proper styling"""
     if role == "user":
@@ -182,25 +135,14 @@ def display_chat_message(role: str, content: str, model: str = "", timestamp: st
 # Main header
 st.markdown("""
 <div class="main-header">
-    <h1>🛡️ SafeInfer LLM Chatbot</h1>
-    <p>Secure and intelligent conversations powered by SafeInfer API</p>
+    <h1>🛡️ Finance Ops Chatbot</h1>
+    <p>Helpful assistant for Finance Ops team</p>
 </div>
 """, unsafe_allow_html=True)
 
 # Sidebar configuration
 with st.sidebar:
-    st.header("⚙️ Configuration")
 
-    # Model selection
-    available_models = AVAILABLE_MODELS
-    if available_models:
-        st.subheader("🤖 Model Selection")
-        selected_model = st.selectbox(
-            "Choose a model:",
-            available_models,
-            index=available_models.index(st.session_state.selected_model)
-        )
-        st.session_state.selected_model = selected_model
 
     # API connection test
     st.subheader("🔗 API Status")
@@ -228,14 +170,27 @@ def display_chat_message(role: str, content: str, model: str = "", timestamp: st
         st.download_button(
             label="📥 Export Chat",
             data=json.dumps(chat_data, indent=2),
-            file_name=f"safe_infer_chat_{time.strftime('%Y%m%d_%H%M%S')}.json",
+            file_name=f"finance_chatbot_{time.strftime('%Y%m%d_%H%M%S')}.json",
             mime="application/json"
         )
 
     # Statistics
     st.subheader("📊 Statistics")
     st.metric("Messages", len(st.session_state.chat_history))
-    st.metric("Current Model", st.session_state.selected_model)
+    st.markdown(f"""
+<div style="font-size:0.8rem;">
+    Current Model: <br><span style="font-size:1.2rem;"><b>{st.session_state.model_name}</b></span>
+</div>
+""", unsafe_allow_html=True)
+
+# Welcome message
+st.markdown(f"""
+<div class="chat-message bot-message">
+    <strong>🤖 AI Assistant:</strong><br>
+    Welcome {USER_EMAIL}. {USER_TEAM} team!
+</div>
+""", unsafe_allow_html=True)
+
 
 # Main chat interface
 st.subheader("💬 Chat Interface")
@@ -261,73 +216,6 @@ def display_chat_message(role: str, content: str, model: str = "", timestamp: st
 col1, col2 = st.columns([1, 4])
 with col1:
     send_button = st.button("🚀 Send", type="primary")
-with col2:
-    regenerate_button = st.button("🔄 Regenerate Last Response")
-    if regenerate_button and st.session_state.chat_history:
-        # Remove the last bot response and regenerate
-        while st.session_state.chat_history and st.session_state.chat_history[-1]["role"] == "assistant":
-            st.session_state.chat_history.pop()
-        if st.session_state.chat_history:
-            # Store the last user message for regeneration
-            last_user_message = st.session_state.chat_history[-1]["content"]
-            # Process the regeneration
-            if last_user_message.strip():
-                # Add user message to history
-                st.session_state.chat_history.append({
-                    "role": "user",
-                    "content": last_user_message,
-                    "timestamp": time.strftime("%H:%M:%S")
-                })
-                
-                # Display user message
-                display_chat_message("user", last_user_message)
-                
-                # Get AI response
-                with st.spinner("🤖 AI is thinking..."):
-                    result = call_safe_infer_api(
-                        message=last_user_message,
-                        model=st.session_state.selected_model,
-                        api_key=st.session_state.api_key
-                    )
-                
-                if result["status"] == "success":
-                    # Extract response content
-                    response_content = extract_response_content(result["data"])
-                    
-                    # Add bot response to history
-                    st.session_state.chat_history.append({
-                        "role": "assistant",
-                        "content": response_content,
-                        "model": st.session_state.selected_model,
-                        "timestamp": time.strftime("%H:%M:%S")
-                    })
-                    
-                    # Display bot response
-                    display_chat_message(
-                        "assistant", 
-                        response_content, 
-                        st.session_state.selected_model,
-                        time.strftime("%H:%M:%S")
-                    )
-                    
-                    # Show classification info if available
-                    if 'response' in result["data"] and isinstance(result["data"]["response"], dict):
-                        response_data = result["data"]["response"]
-                        if 'classification' in response_data:
-                            classification = response_data['classification']
-                            with st.expander("🔍 Response Analysis"):
-                                st.json(classification)
-                
-                else:
-                    error_message = f"❌ Error: {result['message']}"
-                    st.error(error_message)
-                    st.session_state.chat_history.append({
-                        "role": "assistant",
-                        "content": error_message,
-                        "timestamp": time.strftime("%H:%M:%S")
-                    })
-                
-                st.rerun()
 
 # Process user input
 if send_button and user_input.strip():
@@ -343,28 +231,29 @@ def display_chat_message(role: str, content: str, model: str = "", timestamp: st
 
     # Get AI response
     with st.spinner("🤖 AI is thinking..."):
-        result = call_safe_infer_api(
+        model = SELECTED_MODEL
+
+        result = call_open_ai(
             message=user_input,
-            model=st.session_state.selected_model,
+            model=model,
             api_key=st.session_state.api_key
         )
-    
+        result = {"status": "success", "data": result}
     if result["status"] == "success":
         # Extract response content
-        response_content = extract_response_content(result["data"])
-        
-        # Add bot response to history
+        response = result['data']['data']
+
         st.session_state.chat_history.append({
             "role": "assistant",
-            "content": response_content,
+            "content": response,
             "model": st.session_state.selected_model,
             "timestamp": time.strftime("%H:%M:%S")
         })
 
         # Display bot response
         display_chat_message(
             "assistant", 
-            response_content, 
+            response,
             st.session_state.selected_model,
             time.strftime("%H:%M:%S")
         )