Update app.py

app.py

@@ -350,7 +350,7 @@ def load_local_model():
 
 tokenizer, model = load_local_model()
 
-def generate_response(input_dict, use_openai=False):
+def generate_response(input_dict, use_openai=False, max_tokens=700):
     prompt = grantbuddy_prompt.format(**input_dict)
 
     if use_openai:
@@ -419,7 +419,7 @@ def get_rag_chain(retriever, use_openai=False):
         }
 
     return RunnableLambda(merge_contexts) | RunnableLambda(
-        lambda input_dict: generate_response(input_dict, use_openai=use_openai)
+        lambda input_dict: generate_response(input_dict, use_openai=use_openai, max_tokens=max_tokens)
     )
 
 
@@ -428,13 +428,21 @@ def main():
     # st.set_page_config(page_title="Grant Buddy RAG", page_icon="🤖")
     st.title("🤖 Grant Buddy: Grant-Writing Assistant")
     USE_OPENAI = st.sidebar.checkbox("Use OpenAI (Costs Tokens)", value=False)
+    st.sidebar.markdown("### Retrieval Settings")
+
+    k_value = st.sidebar.slider("How many chunks to retrieve (k)", min_value=5, max_value=40, step=5, value=10)
+    score_threshold = st.sidebar.slider("Minimum relevance score", min_value=0.0, max_value=1.0, step=0.05, value=0.75)
+
+    st.sidebar.markdown("### Generation Settings")
+    max_tokens = st.sidebar.number_input("Max tokens in response", min_value=100, max_value=1500, value=700, step=50)
+
     if "generated_queries" not in st.session_state:
         st.session_state.generated_queries = {}
     
     manual_context = st.text_area("📝 Optional: Add your own context (e.g., mission, goals)", height=150)
 
-    retriever = init_vector_search().as_retriever(search_kwargs={"k": 10, "score_threshold": 0.75})
-    rag_chain = get_rag_chain(retriever, use_openai=USE_OPENAI)
+    retriever = init_vector_search().as_retriever(search_kwargs={"k": k_value, "score_threshold": score_threshold})
+    rag_chain = get_rag_chain(retriever, use_openai=USE_OPENAI, max_tokens=max_tokens)
 
     uploaded_file = st.file_uploader("Upload PDF or TXT for extra context (optional)", type=["pdf", "txt"])
     uploaded_text = ""
@@ -475,16 +483,13 @@ def main():
                     for q in selected_questions:
                         # full_query = f"{q}\n\nAdditional context:\n{uploaded_text}"
                         combined_context = "\n\n".join(filter(None, [manual_context.strip(), uploaded_text.strip()]))
-                        response = rag_chain.invoke({
-                                "question": q,
-                                 "manual_context": combined_context
-                            })
-                        # response = rag_chain.invoke(full_query)
-                        # answers.append({"question": q, "answer": response})
                         if q in st.session_state.generated_queries:
                             response = st.session_state.generated_queries[q]
                         else:
-                            response = rag_chain.invoke(full_query)
+                            response = rag_chain.invoke({
+                                "question": q,
+                                "manual_context": combined_context
+                            })
                             st.session_state.generated_queries[q] = response
                         answers.append({"question": q, "answer": response})
                 for item in answers: