Update app.py

app.py

@@ -2,51 +2,80 @@ import os
 import gradio as gr
 from anthropic import Anthropic
 from pypdf import PdfReader
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
 
-# Set up username and password
+# Set up your Anthropic API key in HF secrets
+ANTHROPIC_API_KEY = os.getenv('ANTHROPIC_API_KEY')
+os.environ["ANTHROPIC_API_KEY"] = ANTHROPIC_API_KEY
+
+# Set up username and password in HF secrets
 username = os.getenv('username')
 password = os.getenv('password')
 
 # Add the path to your desired knowledge base
 reference_document = "Rosenshine+Principles+red.pdf" 
 reader = PdfReader(reference_document)
-text = ''.join(page.extract_text() for page in reader.pages)
+full_text = ''.join(page.extract_text() for page in reader.pages)
+text_chunks = chunk_text(full_text)
 
-# Anthropic API setup
-ANTHROPIC_API_KEY = os.getenv('ANTHROPIC_API_KEY')
-os.environ["ANTHROPIC_API_KEY"] = ANTHROPIC_API_KEY
+# Function to chunk the document
+def chunk_text(text, chunk_size=1000, overlap=100):
+    chunks = []
+    start = 0
+    while start < len(text):
+        end = start + chunk_size
+        chunk = text[start:end]
+        chunks.append(chunk)
+        start = end - overlap
+    return chunks
+
+# Function to find the most relevant chunks
+def get_relevant_chunks(query, chunks, top_n=3):
+    vectorizer = TfidfVectorizer()
+    tfidf_matrix = vectorizer.fit_transform(chunks + [query])
+    cosine_similarities = cosine_similarity(tfidf_matrix[-1], tfidf_matrix[:-1]).flatten()
+    relevant_indices = cosine_similarities.argsort()[-top_n:][::-1]
+    return [chunks[i] for i in relevant_indices]
 
 def chat_with_assistant(message, history):
-    history_str = "\n".join([f"Human: {h[0]}\nAssistant: {h[1]}" for h in history])
+    # Find relevant chunks based on the user message
+    relevant_chunks = get_relevant_chunks(message, text_chunks)
+    context = "\n".join(relevant_chunks)
     
+    # Prepare the system message
     ai_message = f"""You are an AI assistant answering questions based on a reference document. 
-    You provide short, clear answers with simple language.
-    Use the following text as context for all of your answers:
-
-    {text}
-
-    Previous conversation history:
-    {history_str}
+    You provide short, clear answers in simple language.
+    Use the following as context for all of your answers:
+    {context}
     """
 
-    # Add your desired instructions
-    
+    # Customize instructions as needed
     instructions = """
     
-    """ 
+    """
+    system message = f"{ai_message} {instructions}"
+                  
+    # Prepare the message array
+    messages = [{"role": "system", "content": system_message}]
+    
+    # Add conversation history
+    for human_msg, ai_msg in history:
+        messages.append({"role": "user", "content": human_msg})
+        messages.append({"role": "assistant", "content": ai_msg})
     
-    system_prompt = f"{ai_message} {instructions}"
-
+    # Add the current user message
+    messages.append({"role": "user", "content": message})
+    
+    # Create Anthropic client
     client = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
     
+    # Make the API call
     response = client.messages.create(
-        model="claude-3-5-sonnet-20240620",
-      # model="claude-3-haiku-20240307",
+        model="claude-3-sonnet-20240307",
+      # model ="claude-3-haiku-20240307", 
         max_tokens=500,
-        system=system_prompt,
-        messages=[
-            {"role": "user", "content": message}
-        ]
+        messages=messages
     )
     
     return response.content[0].text.strip()