Update app.py

app.py

@@ -4,51 +4,86 @@ from bs4 import BeautifulSoup
 import re
 from sentence_transformers import SentenceTransformer, util
 import torch
+from transformers import pipeline
 
-# Initialize the sentence transformer model
+# Initialize the sentence transformer model and summarizer
 @st.cache_resource
-def load_model():
-    return SentenceTransformer('all-MiniLM-L6-v2')
+def load_models():
+    return SentenceTransformer('all-MiniLM-L6-v2'), pipeline("summarization")
 
-model = load_model()
+model, summarizer = load_models()
 
 # Function to scrape documentation
 def scrape_documentation(url):
     response = requests.get(url)
     soup = BeautifulSoup(response.text, 'html.parser')
-    # This pattern might need adjustment based on the specific structure of the documentation
-    code_blocks = soup.find_all('pre')
-    return [block.text for block in code_blocks]
+    sections = soup.find_all(['h1', 'h2', 'h3', 'p', 'pre'])
+    
+    content = []
+    current_section = {"title": "", "content": "", "code": ""}
+    
+    for section in sections:
+        if section.name in ['h1', 'h2', 'h3']:
+            if current_section["title"]:
+                content.append(current_section)
+            current_section = {"title": section.text.strip(), "content": "", "code": ""}
+        elif section.name == 'p':
+            current_section["content"] += section.text.strip() + " "
+        elif section.name == 'pre':
+            current_section["code"] += section.text.strip() + "\n"
+    
+    if current_section["title"]:
+        content.append(current_section)
+    
+    return content
 
-# Function to find the most relevant code snippet
-def find_relevant_snippet(query, snippets):
+# Function to find the most relevant section
+def find_relevant_section(query, sections):
     query_embedding = model.encode(query, convert_to_tensor=True)
-    snippet_embeddings = model.encode(snippets, convert_to_tensor=True)
+    section_embeddings = model.encode([s["title"] + " " + s["content"] for s in sections], convert_to_tensor=True)
     
-    cosine_scores = util.pytorch_cos_sim(query_embedding, snippet_embeddings)
+    cosine_scores = util.pytorch_cos_sim(query_embedding, section_embeddings)
     best_match = torch.argmax(cosine_scores)
     
-    return snippets[best_match]
+    return sections[best_match]
+
+# Function to summarize text
+def summarize_text(text, max_length=150):
+    return summarizer(text, max_length=max_length, min_length=30, do_sample=False)[0]['summary_text']
 
 # Streamlit UI
-st.title("AI Code Assistant for Library Documentation")
+st.title("Enhanced AI Code Assistant for Library Documentation")
 
 library = st.selectbox("Choose a library", ["llama-index", "langchain"])
 
 if library == "llama-index":
-    url = "https://gpt-index.readthedocs.io/en/latest/examples/index.html"
+    url = "https://gpt-index.readthedocs.io/en/latest/getting_started/installation.html"
 elif library == "langchain":
     url = "https://python.langchain.com/docs/get_started/introduction"
 
-query = st.text_input("What kind of code snippet are you looking for?")
+query = st.text_input("What would you like to know about the library?")
 
-if st.button("Find Code Snippet"):
-    with st.spinner("Searching for relevant code..."):
-        snippets = scrape_documentation(url)
-        if snippets:
-            relevant_snippet = find_relevant_snippet(query, snippets)
-            st.code(relevant_snippet, language="python")
+if st.button("Get Information"):
+    with st.spinner("Searching and processing information..."):
+        sections = scrape_documentation(url)
+        if sections:
+            relevant_section = find_relevant_section(query, sections)
+            
+            st.subheader(relevant_section["title"])
+            
+            if relevant_section["content"]:
+                summary = summarize_text(relevant_section["content"])
+                st.write("Summary:", summary)
+            
+            if relevant_section["code"]:
+                st.subheader("Code Example:")
+                st.code(relevant_section["code"], language="python")
+                
+                code_summary = summarize_text(f"This code {relevant_section['code']}")
+                st.write("Code summary:", code_summary)
+            
+            st.write("For more detailed information, please refer to the official documentation.")
         else:
-            st.error("No code snippets found in the documentation.")
+            st.error("No relevant information found in the documentation.")
 
-st.warning("Note: This tool scrapes the latest documentation, but may not always return perfect results. Always verify the code and check the official documentation.")
+st.warning("Note: This tool provides information based on the latest documentation, but may not always return perfect results. Always verify the information and check the official documentation for the most up-to-date and accurate details.")