Update app.py

app.py

@@ -1,69 +1,84 @@
 import streamlit as st
 import os
-from langchain_huggingface import HuggingFaceEndpoint
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
+from langchain_core.messages import HumanMessage
 
 # --- 1. UI Setup ---
-st.set_page_config(page_title="Prompting Demo", page_icon="🤖")
-st.title("Prompt Engineering Demo")
-st.markdown("Experiment with Zero-Shot, Few-Shot, and Chain of Thought.")
+st.set_page_config(page_title="ShotCaller AI", page_icon="🤖")
+st.title("🤖 ShotCaller: Prompt Lab")
+st.markdown("Explore how different prompting techniques change AI reasoning.")
 
 # --- 2. Model Setup ---
-# On HF Spaces, you store keys in "Settings > Variables and Secrets"
+# Retrieve your token from Hugging Face Settings > Secrets
 api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
 if not api_token:
-    st.error("Please add your HUGGINGFACEHUB_API_TOKEN to Secrets.")
+    st.error("Missing HUGGINGFACEHUB_API_TOKEN. Add it to your Space Secrets!")
     st.stop()
 
+# Initialize the endpoint
+# Note: task="text-generation" is required for the underlying class, 
+# but ChatHuggingFace will handle the conversational routing.
 repo_id = "mistralai/Mistral-7B-Instruct-v0.3"
 llm = HuggingFaceEndpoint(
     repo_id=repo_id,
-    # Remove 'task' and let HF auto-detect
-    huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"]
+    task="text-generation",
+    temperature=0.7,
+    huggingfacehub_api_token=api_token
 )
 
+# Wrap the LLM in ChatHuggingFace to satisfy the 'conversational' requirement
+chat_model = ChatHuggingFace(llm=llm)
+
 # --- 3. Sidebar Selection ---
 option = st.sidebar.selectbox(
-    "Choose Prompting Technique",
+    "Select Prompting Strategy",
     ("Zero-Shot", "Single-Shot", "Few-Shot", "Chain of Thought")
 )
 
-user_input = st.text_input("Enter your query:", "How does a solar panel work?")
+user_query = st.text_input("What's your question?", "Why is the sky blue?")
+
+# --- 4. Logic for Prompting Techniques ---
+if st.button("Run Inference"):
+    formatted_prompt = ""
 
-# --- 4. Logic for each Technique ---
-if st.button("Generate Response"):
     if option == "Zero-Shot":
-        prompt = PromptTemplate.from_template("{input}")
-        result = llm.invoke(prompt.format(input=user_input))
+        # Direct question, no guidance
+        formatted_prompt = user_query
 
-    elif option == "Single-Shot" or option == "Few-Shot":
-        # Examples for sentiment analysis or translation
+    elif option == "Single-Shot":
+        # One example to set the tone/format
+        formatted_prompt = f"Example: Input: Hello, Output: Hi there!\nInput: {user_query}\nOutput:"
+
+    elif option == "Few-Shot":
+        # Multiple examples for pattern recognition
         examples = [
-            {"input": "I am happy", "output": "Positive"},
-            {"input": "This is bad", "output": "Negative"}
+            {"input": "The movie was great", "output": "Positive"},
+            {"input": "The food was cold", "output": "Negative"},
+            {"input": "It was an okay experience", "output": "Neutral"}
         ]
-        # Use only one example for single-shot
-        ex_to_use = examples[:1] if option == "Single-Shot" else examples
-        
         example_prompt = PromptTemplate(input_variables=["input", "output"], template="Input: {input}\nOutput: {output}")
-        few_shot_prompt = FewShotPromptTemplate(
-            examples=ex_to_use,
+        few_shot_p = FewShotPromptTemplate(
+            examples=examples,
             example_prompt=example_prompt,
             suffix="Input: {input}\nOutput:",
             input_variables=["input"]
         )
-        result = llm.invoke(few_shot_prompt.format(input=user_input))
+        formatted_prompt = few_shot_p.format(input=user_query)
 
     elif option == "Chain of Thought":
-        cot_template = """
-        Question: If John has 5 pears and gives 2 to Mary, how many does he have?
-        Answer: Let's think step by step. John starts with 5. He gives 2 away. 5 - 2 = 3. The answer is 3.
-
-        Question: {input}
-        Answer: Let's think step by step."""
-        prompt = PromptTemplate.from_template(cot_template)
-        result = llm.invoke(prompt.format(input=user_input))
+        # Encouraging step-by-step logic
+        formatted_prompt = f"Question: {user_query}\nAnswer: Let's think step by step."
 
-    st.subheader(f"Result ({option})")
-    st.write(result)
+    # Execute the call
+    with st.spinner("Thinking..."):
+        try:
+            # Wrap our formatted string in a HumanMessage for the Chat Model
+            messages = [HumanMessage(content=formatted_prompt)]
+            response = chat_model.invoke(messages)
+            
+            st.subheader(f"Result: {option}")
+            st.write(response.content)
+        except Exception as e:
+            st.error(f"An error occurred: {e}")