use openai oss

src/Rag.py

@@ -7,6 +7,7 @@ import numpy as np
 from dotenv import load_dotenv
 from sentence_transformers import SentenceTransformer
 from together import Together
+from openai import OpenAI
 
 global db, referenced_tables_db, embedder, index, llm_client
 
@@ -121,6 +122,14 @@ def load_together_llm_client():
     load_dotenv()
     return Together(api_key=os.getenv("TOGETHER_API_KEY"))
 
+def load_nvidia_llm_client():
+    load_dotenv()
+    return OpenAI(
+        base_url="https://integrate.api.nvidia.com/v1",
+        api_key=os.getenv("NVIDIA_API_KEY"),
+    )
+
+
 
 # -------- Prompt Construction --------
 def construct_prompt(query, faiss_results):
@@ -193,6 +202,40 @@ def call_llm(llm_client, prompt, stream_flag=False, max_tokens=500, temperature=
         traceback.print_exc()
         raise
 
+def call_nvidia_llm(llm_client, prompt, stream_flag=False, max_tokens=4096, temperature=0.6, top_p=0.7, model_name="openai/gpt-oss-20b"):
+    print(f"Calling NVIDIA LLM with model: {model_name}")
+    try:
+        if stream_flag:
+            def stream_generator():
+                completion = llm_client.chat.completions.create(
+                    model=model_name,
+                    messages=[{"role":"user","content": prompt}],
+                    temperature=temperature,
+                    top_p=top_p,
+                    max_tokens=max_tokens,
+                    stream=True
+                )
+                for chunk in completion:
+                    if chunk.choices[0].delta.content is not None:
+                        yield chunk.choices[0].delta.content
+            return stream_generator()
+        else:
+            completion = llm_client.chat.completions.create(
+                model=model_name,
+                messages=[{"role":"user","content": prompt}],
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False
+            )
+            return completion.choices[0].message.content
+    except Exception as e:
+        print("Error in call_nvidia_llm:", str(e))
+        import traceback
+        traceback.print_exc()
+        raise
+
+
 
 def call_ollama(prompt, model="mistral", stream_flag=False, max_tokens=500, temperature=0.05, top_p=0.9):
     url = "http://localhost:11434/api/generate"
@@ -226,25 +269,57 @@ def launch_depression_assistant(embedder_name, designated_client=None):
     index = load_cosine_index(embedder_name)
 
     if designated_client is None:
-        print("No LLM client provided. Loading Together LLM client...")
+        print("Attempting to load NVIDIA LLM client...")
         try:
-            llm_client = load_together_llm_client()
-        except Exception:
-            print("Failed to load Together LLM client. Please check your API key.")
+            llm_client = load_nvidia_llm_client()
+            print("Successfully loaded NVIDIA LLM client.")
+        except Exception as e:
+            print(f"Failed to load NVIDIA LLM client: {e}")
+            print("Attempting to load Together LLM client as a fallback...")
+            try:
+                llm_client = load_together_llm_client()
+                print("Successfully loaded Together LLM client.")
+            except Exception as e:
+                print(f"Failed to load Together LLM client: {e}")
+                llm_client = None
     else:
         llm_client = designated_client
+        print(f"Using designated client: {type(llm_client).__name__}")
+
+    if llm_client is None:
+        print("Warning: No LLM client could be loaded. The assistant will not be able to generate responses.")
 
     print("---------Depression Assistant is ready to use!--------------\n\n")
 
 
-def depression_assistant(query, model_name="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free", max_tokens=500, temperature=0.05, top_p=0.9, stream_flag=False, chat_history=None):
+
+def depression_assistant(query, model_name=None, max_tokens=None, temperature=None, top_p=None, stream_flag=False, chat_history=None):
     results = vector_search(query, embedder, db, index, referenced_tables_db, k=3)
     prompt = construct_prompt_with_memory(query, results, chat_history=chat_history)
 
+    kwargs = {}
+    if model_name:
+        kwargs['model_name'] = model_name
+    if max_tokens:
+        kwargs['max_tokens'] = max_tokens
+    if temperature is not None:
+        kwargs['temperature'] = temperature
+    if top_p:
+        kwargs['top_p'] = top_p
+
     if llm_client == "Run Ollama Locally":
-        return results, call_ollama(prompt, model_name, stream_flag, max_tokens, temperature, top_p)
+        if 'model_name' in kwargs:
+            kwargs['model'] = kwargs.pop('model_name')
+        return results, call_ollama(prompt, stream_flag=stream_flag, **kwargs)
+    elif isinstance(llm_client, OpenAI):  # NVIDIA Client
+        return results, call_nvidia_llm(llm_client, prompt, stream_flag=stream_flag, **kwargs)
+    elif isinstance(llm_client, Together):  # Together Client
+        return results, call_llm(llm_client, prompt, stream_flag=stream_flag, **kwargs)
     else:
-        return results, call_llm(llm_client, prompt, stream_flag, max_tokens, temperature, top_p, model_name)
+        if llm_client is None:
+            raise ValueError("LLM client not initialized. Please check API keys.")
+        # Fallback to NVIDIA as requested
+        return results, call_nvidia_llm(llm_client, prompt, stream_flag=stream_flag, **kwargs)
 
 
 def load_queries_and_answers(query_file, answers_file):

src/app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from Rag import launch_depression_assistant, depression_assistant
+import Rag
 from openai import OpenAI
 from together import Together
 import time
@@ -47,32 +47,33 @@ with st.sidebar:
     if embedder_name == "Other":
         embedder_name = st.text_input('Enter the embedder model name')
 
-    llm_client = Together(api_key=os.getenv("TOGETHER_API_KEY"))
-
     if (not st.session_state.embedder_loaded or
             st.session_state.current_embedder_name != embedder_name):
 
         with st.spinner(f"Loading embedding model: {embedder_name}..."):
-            launch_depression_assistant(embedder_name=embedder_name, designated_client=llm_client)
+            Rag.launch_depression_assistant(embedder_name=embedder_name)
             st.session_state.embedder_loaded = True
             st.session_state.current_embedder_name = embedder_name
             st.success(f"✅ Embedding model {embedder_name} loaded successfully!")
     else:
         st.info(f"📋 Current embedding model: {st.session_state.current_embedder_name}")
 
+    if isinstance(Rag.llm_client, OpenAI):
+        # NVIDIA client
+        model_list = ["openai/gpt-oss-20b"]
+    elif isinstance(Rag.llm_client, Together):
+        # Together client
+        model_list = ["meta-llama/Llama-3.3-70B-Instruct-Turbo-Free",
+                      "deepseek-ai/deepseek-r1",
+                      "meta/llama-3.3-70b-instruct"]
+    else:
+        # Default or unknown client
+        model_list = ["meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"]
+
     selected_model = st.selectbox('Choose a model for generation',
-                                  ["meta-llama/Llama-3.3-70B-Instruct-Turbo-Free",
-                                   "deepseek-ai/deepseek-r1",
-                                   "meta/llama-3.3-70b-instruct"],
+                                  model_list,
                                   key='selected_model')
 
-    if selected_model in ["deepseek-ai/deepseek-r1", "meta/llama-3.3-70b-instruct"]:
-        max_length = 1000
-        llm_client = OpenAI(
-            base_url="https://integrate.api.nvidia.com/v1",
-            api_key=os.getenv("NVIDIA_API_KEY", None),
-        )
-
     temperature = st.slider('temperature', min_value=0.01, max_value=1.0, value=0.05, step=0.01)
     top_p = st.slider('top_p', min_value=0.01, max_value=1.0, value=0.9, step=0.01)
     max_length = st.slider('max_length', min_value=100, max_value=1000, value=500, step=10)
@@ -145,7 +146,7 @@ with chat_col:
 
         collected = ""
         t0 = time.perf_counter()
-        results, response = depression_assistant(user_input, model_name=selected_model, max_tokens=max_length,
+        results, response = Rag.depression_assistant(user_input, model_name=selected_model, max_tokens=max_length,
                                                  temperature=temperature, top_p=top_p, stream_flag=True,
                                                  chat_history=history)