Update app.py

app.py

@@ -12,7 +12,7 @@ from pymongo import MongoClient
 from PyPDF2 import PdfReader
 st.set_page_config(page_title="Grant Buddy RAG", page_icon="🤖")
 
-
+from typing import List
 
 from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
 from langchain.embeddings import HuggingFaceEmbeddings
@@ -229,9 +229,7 @@ def init_vector_search() -> MongoDBAtlasVectorSearch:
 #         if len(clean) > 10 and not any(bad in clean.lower() for bad in ["phone", "email", "address", "website"]):
 #             prompts.append(clean)
 #     return prompts
-from typing import List
-import os
-import openai
+
 
 def extract_with_llm_local(text: str, use_openai: bool = False) -> List[str]:
     # Example context to prime the model
@@ -266,7 +264,7 @@ PROMPTS:
             return "⚠️ OpenAI key missing."
         try:
             response = client.chat.completions.create(
-                model="gpt-3.5-turbo",
+                model="gpt-4o-mini",
                 messages=[
                     {"role": "system", "content": "You extract prompts and headers from grant text."},
                     {"role": "user", "content": prompt},
@@ -276,6 +274,8 @@ PROMPTS:
             )
             # raw_output = response["choices"][0]["message"]["content"]
             raw_output = response.choices[0].message.content
+            st.markdown(f"🧮 Extract Tokens: Prompt = {response.usage.prompt_tokens}, "
+                f"Completion = {response.usage.completion_tokens}, Total = {response.usage.total_tokens}")
         except Exception as e:
             st.error(f"❌ OpenAI extraction failed: {e}")
             return []
@@ -351,16 +351,12 @@ def load_local_model():
 tokenizer, model = load_local_model()
 
 def generate_response(input_dict, use_openai=False):
-    if use_openai:
-        if not openai.api_key:
-            st.error("❌ OPENAI_API_KEY is not set.")
-            return "⚠️ OpenAI key missing."
-
-        prompt = grantbuddy_prompt.format(**input_dict)
+    prompt = grantbuddy_prompt.format(**input_dict)
 
+    if use_openai:
         try:
             response = client.chat.completions.create(
-                model="gpt-3.5-turbo", 
+                model="gpt-4o-mini", 
                 messages=[
                     {"role": "system", "content": prompt},
                     {"role": "user", "content": input_dict["question"]},
@@ -368,14 +364,30 @@ def generate_response(input_dict, use_openai=False):
                 temperature=0.2,
                 max_tokens=700,
             )
-            return response.choices[0].message.content.strip()
+            answer = response.choices[0].message.content.strip()
+
+            # ✅ Token logging
+            prompt_tokens = response.usage.prompt_tokens
+            completion_tokens = response.usage.completion_tokens
+            total_tokens = response.usage.total_tokens
+
+            return {
+                "answer": answer,
+                "tokens": {
+                    "prompt": prompt_tokens,
+                    "completion": completion_tokens,
+                    "total": total_tokens
+                }
+            }
+
         except Exception as e:
             st.error(f"❌ OpenAI error: {e}")
-            return "⚠️ OpenAI request failed."
+            return {
+                "answer": "⚠️ OpenAI request failed.",
+                "tokens": {}
+            }
 
     else:
-        # Local TinyLlama path
-        prompt = grantbuddy_prompt.format(**input_dict)
         inputs = tokenizer(prompt, return_tensors="pt")
         outputs = model.generate(
             **inputs,
@@ -385,17 +397,31 @@ def generate_response(input_dict, use_openai=False):
             pad_token_id=tokenizer.eos_token_id
         )
         decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return decoded[len(prompt):].strip()
+        return {
+            "answer": decoded[len(prompt):].strip(),
+            "tokens": {}
+        }
 
 
 
 
 # =================== RAG Chain ===================
 def get_rag_chain(retriever, use_openai=False):
-    return {
-        "context": retriever | RunnableLambda(format_docs),
-        "question": RunnablePassthrough()
-    } | RunnableLambda(lambda input_dict: generate_response(input_dict, use_openai=use_openai))
+    def merge_contexts(inputs):
+        retrieved_chunks = format_docs(retriever.invoke(inputs["question"]))
+        combined = "\n\n".join(filter(None, [
+            inputs.get("manual_context", ""),
+            retrieved_chunks
+        ]))
+        return {
+            "context": combined,
+            "question": inputs["question"]
+        }
+
+    return RunnableLambda(merge_contexts) | RunnableLambda(
+        lambda input_dict: generate_response(input_dict, use_openai=use_openai)
+    )
+
 
 # =================== Streamlit UI ===================
 def main():
@@ -404,7 +430,8 @@ def main():
     USE_OPENAI = st.sidebar.checkbox("Use OpenAI (Costs Tokens)", value=False)
     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)
@@ -440,12 +467,18 @@ def main():
                     selected_questions.append(q)
             submit_button = st.form_submit_button("Submit")
             
+        #Multi-Select Question
         if 'submit_button' in locals() and submit_button:
             if selected_questions:
                 with st.spinner("💡 Generating answers..."):
                     answers = []
                     for q in selected_questions:
-                        full_query = f"{q}\n\nAdditional context:\n{uploaded_text}"
+                        # 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:
@@ -456,29 +489,16 @@ def main():
                         answers.append({"question": q, "answer": response})
                 for item in answers:
                     st.markdown(f"### ❓ {item['question']}")
-                    st.markdown(f"💬 {item['answer']}")
+                    st.markdown(f"💬 {item['answer']['answer']}")
+                    tokens = item['answer'].get("tokens", {})
+                    if tokens:
+                        st.markdown(f"🧮 **Token Usage:** Prompt = {tokens.get('prompt')}, "
+                        f"Completion = {tokens.get('completion')}, Total = {tokens.get('total')}")
+                    
             else:
                 st.info("No prompts selected for answering.")
 
             
-
-        # #select prompts to answer
-        # selected_questions = st.multiselect("✅ Choose prompts to answer:", filtered_questions, default=filtered_questions)
-
-        # if selected_questions:
-        #     with st.spinner("💡 Generating answers..."):
-        #         answers = []
-        #         for q in selected_questions:
-        #             full_query = f"{q}\n\nAdditional context:\n{uploaded_text}"
-        #             response = rag_chain.invoke(full_query)
-        #             answers.append({"question": q, "answer": response})
-
-        #     for item in answers:
-        #         st.markdown(f"### ❓ {item['question']}")
-        #         st.markdown(f"💬 {item['answer']}")
-        # else:
-        #     st.info("No prompts selected for answering.")
-
     # ✍️ Manual single-question input
     query = st.text_input("Ask a grant-related question")
     if st.button("Submit"):
@@ -486,13 +506,19 @@ def main():
             st.warning("Please enter a question.")
             return
 
-        full_query = f"{query}\n\nAdditional context:\n{uploaded_text}" if uploaded_text else query
+        # full_query = f"{query}\n\nAdditional context:\n{uploaded_text}" if uploaded_text else query
+        combined_context = "\n\n".join(filter(None, [manual_context.strip(), uploaded_text.strip()]))
         with st.spinner("🤖 Thinking..."):
-            response = rag_chain.invoke(full_query)
-            st.text_area("Grant Buddy says:", value=response, height=250, disabled=True)
-
+            # response = rag_chain.invoke(full_query)
+            response = rag_chain.invoke({"question":query,"manual_context": combined_context})
+            st.text_area("Grant Buddy says:", value=response["answer"], height=250, disabled=True)
+            tokens=response.get("tokens",{})
+            if tokens:
+                 st.markdown(f"🧮 **Token Usage:** Prompt = {tokens.get('prompt')}, "
+                 f"Completion = {tokens.get('completion')}, Total = {tokens.get('total')}")
+                
         with st.expander("🔍 Retrieved Chunks"):
-            context_docs = retriever.get_relevant_documents(full_query)
+            context_docs = retriever.get_relevant_documents(query)
             for doc in context_docs:
                 # st.json(doc.metadata)
                 st.markdown(f"**Chunk ID:** {doc.metadata.get('chunk_id', 'unknown')} | **Title:** {doc.metadata['metadata'].get('title', 'unknown')}")