Update app.py

app.py

@@ -1,39 +1,27 @@
-import os
-import requests
-import numpy as np
-import faiss
+import os, requests, numpy as np, faiss
 import gradio as gr
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import InferenceClient
 
-# 1. INITIALIZE MODELS
+# 1. MODELS
 embedder = SentenceTransformer('all-MiniLM-L6-v2')
-# Ensure you have HF_TOKEN in your Space Secrets!
 client = InferenceClient("meta-llama/Llama-3.3-70B-Instruct", token=os.getenv("HF_TOKEN"))
 
 def fetch_and_index(query):
-    """Fetches live book data and indexes it in FAISS."""
     try:
-        url = f"https://openlibrary.org/search.json?q={query}&limit=8"
-        response = requests.get(url, timeout=5)
-        data = response.json().get("docs", [])
-        
-        if not data: return None, None
-
-        catalog = [
-            f"{d.get('title')} by {', '.join(d.get('author_name', ['Unknown Author']))}" 
-            for d in data
-        ]
+        url = f"https://openlibrary.org/search.json?q={query}&limit=5"
+        data = requests.get(url, timeout=5).json().get("docs", [])
+        if not data: return None
+        catalog = [f"{d.get('title')} by {', '.join(d.get('author_name', ['Unknown']))}" for d in data]
         
         embeddings = embedder.encode(catalog)
         index = faiss.IndexFlatL2(embeddings.shape[1])
         index.add(np.array(embeddings).astype('float32'))
-        
         return index, catalog
-    except Exception:
-        return None, None
+    except: return None, None
 
-def librarian_logic(message, history, user_state):
+def librarian_bot(message, history, user_state):
+    # Initialize State
     if user_state is None:
         user_state = {"step": "ASK_AGE", "age": None, "location": None}
 
@@ -42,47 +30,47 @@ def librarian_logic(message, history, user_state):
         if message.isdigit():
             user_state["age"] = int(message)
             user_state["step"] = "ASK_LOCATION"
-            reply = "Got it. For safety compliance, what is your general location (City/Country)?"
-            return history + [[message, reply]], user_state, ""
-        return history + [[message, "How old are you?"]], user_state, ""
+            reply = "Got it. What is your City/Country for compliance?"
+        else:
+            reply = "Welcome! I am your Librarian. To start, how old are you?"
+        return history + [[message, reply]], user_state
 
     if user_state["step"] == "ASK_LOCATION":
         user_state["location"] = message
         user_state["step"] = "SEARCH_READY"
-        reply = f"System verified for {user_state['location']}. What book can I find for you?"
-        return history + [[message, reply]], user_state, ""
+        reply = f"Verification complete. I'm searching for books in {message}. What can I find for you?"
+        return history + [[message, reply]], user_state
 
     # --- PHASE 2: SEARCH & LLM ---
     index, catalog = fetch_and_index(message)
     context = "\n- ".join(catalog[:3]) if catalog else "No books found."
-
-    # THE FIX: This specific list-of-dicts format
+    
     messages = [
-        {"role": "system", "content": "You are a professional librarian. Respond only with book titles and authors based on the provided context."},
-        {"role": "user", "content": f"User Age: {user_state['age']}. Context: {context}. Query: {message}"}
+        {"role": "system", "content": "You are a professional librarian. Give titles and authors."},
+        {"role": "user", "content": f"Context: {context}\nQuery: {message}"}
     ]
 
+    # Generating response
     response = ""
-    # Explicitly use messages=messages
-    for msg in client.chat_completion(
-        messages=messages, 
-        max_tokens=300, 
-        stream=True
-    ):
-        token = msg.choices[0].delta.content
-        if token:
-            response += token
-            yield history + [[message, response]], user_state, ""
+    try:
+        # Non-streaming call for maximum reliability in this debug phase
+        completion = client.chat_completion(messages=messages, max_tokens=200)
+        response = completion.choices[0].message.content
+    except Exception as e:
+        response = f"Connection error: {str(e)}"
 
-# --- UI SETUP ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 📚 AI Librarian Agent\n*Search for books by title or author.*")
-    
-    user_state = gr.State()
+    return history + [[message, response]], user_state
+
+# --- CLEAN UI SETUP ---
+with gr.Blocks() as demo:
+    gr.Markdown("# 📚 AI Librarian Agent")
+    state = gr.State()
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your Input", placeholder="Type age first...")
+    msg = gr.Textbox(placeholder="Enter age, then location, then book query...")
     
-    msg.submit(librarian_logic, [msg, chatbot, user_state], [chatbot, user_state, msg])
+    # We use a simple .submit() that updates the chatbot and state simultaneously
+    msg.submit(librarian_bot, [msg, chatbot, state], [chatbot, state])
+    # Clear the textbox after submit
+    msg.submit(lambda: "", None, [msg])
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()