Update main.py

main.py

@@ -1,85 +1,140 @@
 import os
-from fastapi import FastAPI, HTTPException
-import requests
-from bs4 import BeautifulSoup
+import asyncio
+from fastapi import FastAPI, HTTPException, Query
+from dotenv import load_dotenv
 import aiohttp
+from bs4 import BeautifulSoup
 
 # --- Configuration ---
-# It's recommended to use environment variables for sensitive data like API keys.
-# Replace with your actual API key and endpoint.
-LLM_API_URL = os.getenv("LLM_API_URL", "https://api.inference.net/v1/chat/completions")
-LLM_API_KEY = os.getenv("LLM_API_KEY", "inference-00050468cc1c4a20bd5ca0997c752329") # Replace with your key
+load_dotenv()
+LLM_API_KEY = os.getenv("LLM_API_KEY")
+
+if not LLM_API_KEY:
+    raise RuntimeError("LLM_API_KEY must be set in a .env file.")
+
+# Snapzion Search API Configuration
+SNAPZION_API_URL = "https://search.snapzion.com/get-snippets"
+SNAPZION_HEADERS = {
+    'accept': '*/*',
+    'accept-language': 'en-US,en;q=0.9',
+    'content-type': 'application/json',
+    'origin': 'https://search.snapzion.com',
+    'priority': 'u=1, i',
+    'referer': 'https://search.snapzion.com/docs',
+    'sec-ch-ua': '"Chromium";v="140", "Not=A?Brand";v="24", "Google Chrome";v="140"',
+    'sec-ch-ua-mobile': '?0',
+    'sec-ch-ua-platform': '"Windows"',
+    'sec-fetch-dest': 'empty',
+    'sec-fetch-mode': 'cors',
+    'sec-fetch-site': 'same-origin',
+    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36',
+}
+
+# LLM Configuration
+LLM_API_URL = "https://api.inference.net/v1/chat/completions"
 LLM_MODEL = "meta-llama/llama-3.1-8b-instruct/fp-8"
 
+# --- FastAPI App Initialization ---
 app = FastAPI(
-    title="Web Scraper and AI Processor",
-    description="An API to scrape web content and process it with a large language model.",
-    version="1.0.0"
+    title="AI Search Snippets API (Snapzion)",
+    description="Provides AI-generated summaries from Snapzion search results.",
+    version="1.0.1"
 )
 
-async def scrape_url(session, url: str):
-    """Asynchronously scrapes the text content from a given URL."""
+# --- Core Asynchronous Functions ---
+
+async def call_snapzion_search(session: aiohttp.ClientSession, query: str) -> list:
+    """Calls the Snapzion search API and returns a list of organic results."""
     try:
-        async with session.get(url, timeout=10) as response:
+        async with session.post(SNAPZION_API_URL, headers=SNAPZION_HEADERS, json={"query": query}, timeout=15) as response:
             response.raise_for_status()
-            html_content = await response.text()
-            soup = BeautifulSoup(html_content, "html.parser")
-            # Remove script and style elements
-            for script_or_style in soup(["script", "style"]):
-                script_or_style.decompose()
-            # Get text and clean it up
-            text = soup.get_text()
-            lines = (line.strip() for line in text.splitlines())
-            chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
-            return " ".join(chunk for chunk in chunks if chunk)
-    except requests.exceptions.RequestException as e:
-        raise HTTPException(status_code=400, detail=f"Error fetching the URL: {e}")
-
-async def process_with_llm(session, content: str, query: str):
-    """Sends the scraped content and a query to the LLM for processing."""
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {LLM_API_KEY}",
-    }
-    data = {
-        "messages": [
-            {
-                "role": "system",
-                "content": "You are a helpful assistant that analyzes web content."
-            },
-            {
-                "role": "user",
-                "content": f"Based on the following content, please answer this question: '{query}'\n\nContent:\n{content}"
-            }
-        ],
-        "model": LLM_MODEL,
-        "stream": False  # Set to False for a single response
-    }
+            data = await response.json()
+            return data.get("organic_results", [])
+    except Exception as e:
+        raise HTTPException(status_code=503, detail=f"Search service (Snapzion) failed: {e}")
+
+async def scrape_url(session: aiohttp.ClientSession, url: str) -> str:
+    """Asynchronously scrapes the primary text content from a URL, ignoring PDFs."""
+    if url.lower().endswith('.pdf'):
+        return "Content is a PDF, which cannot be scraped."
     try:
-        async with session.post(LLM_API_URL, headers=headers, json=data, timeout=30) as response:
-            response.raise_for_status()
-            return await response.json()
-    except aiohttp.ClientError as e:
-        raise HTTPException(status_code=500, detail=f"Error communicating with the LLM API: {e}")
+        async with session.get(url, timeout=10) as response:
+            if response.status != 200:
+                return f"Error: Failed to fetch {url} with status {response.status}"
+            html = await response.text()
+            soup = BeautifulSoup(html, "html.parser")
+            for tag in soup(['script', 'style', 'nav', 'footer', 'header', 'aside']):
+                tag.decompose()
+            return " ".join(soup.stripped_strings)
+    except Exception as e:
+        return f"Error: Could not scrape {url}. Reason: {e}"
+
+async def get_ai_snippet(query: str, context: str, sources: list) -> str:
+    """Generates a synthesized answer using an LLM based on the provided context."""
+    headers = {"Authorization": f"Bearer {LLM_API_KEY}", "Content-Type": "application/json"}
+    source_list_str = "\n".join([f"[{i+1}] {source['title']}: {source['link']}" for i, source in enumerate(sources)])
+    
+    prompt = f"""
+Based *only* on the provided context from web pages, provide a concise, factual answer to the user's query. Cite every sentence with the corresponding source number(s), like `[1]`, `[2]`, or `[1, 3]`.
+
+Sources:
+{source_list_str}
 
-@app.post("/scrape-and-process/")
-async def scrape_and_process(url: str, query: str):
+Context:
+---
+{context}
+---
+
+User Query: "{query}"
+
+Answer with citations:
+"""
+    data = {"model": LLM_MODEL, "messages": [{"role": "user", "content": prompt}], "max_tokens": 500}
+
+    async with aiohttp.ClientSession() as session:
+        try:
+            async with session.post(LLM_API_URL, headers=headers, json=data, timeout=45) as response:
+                response.raise_for_status()
+                result = await response.json()
+                return result['choices'][0]['message']['content']
+        except Exception as e:
+            raise HTTPException(status_code=502, detail=f"Failed to get response from LLM: {e}")
+
+# --- API Endpoint ---
+
+@app.get("/search")
+async def ai_search(q: str = Query(..., min_length=3, description="The search query.")):
     """
-    Scrapes a URL, sends the content to a large language model with a query,
-    and returns the model's response.
+    Performs an AI-powered search using Snapzion. It finds relevant web pages,
+    scrapes their content, and generates a synthesized answer with citations.
     """
     async with aiohttp.ClientSession() as session:
-        scraped_content = await scrape_url(session, url)
-        if not scraped_content:
-            raise HTTPException(status_code=404, detail="Could not scrape any content from the URL.")
+        # 1. Search for relevant web pages using Snapzion
+        search_results = await call_snapzion_search(session, q)
+        if not search_results:
+            raise HTTPException(status_code=404, detail="Could not find any relevant sources for the query.")
+        
+        # Limit to the top 4 results for speed and relevance
+        sources = search_results[:4]
 
-        llm_response = await process_with_llm(session, scraped_content, query)
-        return llm_response
+        # 2. Scrape all pages concurrently for speed
+        scrape_tasks = [scrape_url(session, source["link"]) for source in sources]
+        scraped_contents = await asyncio.gather(*scrape_tasks)
 
-@app.get("/")
-def read_root():
-    return {"message": "Welcome to the Web Scraper and AI Processor API."}
+        # 3. Combine content and snippets for a rich context
+        full_context = "\n\n".join(
+            f"Source [{i+1}] (from {sources[i]['link']}):\nOriginal Snippet: {sources[i]['snippet']}\nScraped Content: {content}"
+            for i, content in enumerate(scraped_contents) if not content.startswith("Error:")
+        )
+
+        if not full_context.strip():
+            raise HTTPException(status_code=500, detail="Failed to scrape content from all available sources.")
+
+        # 4. Generate the final AI snippet
+        ai_summary = await get_ai_snippet(q, full_context, sources)
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    return {"ai_summary": ai_summary, "sources": sources}
+
+@app.get("/")
+def root():
+    return {"message": "AI Search API is active. Use the /docs endpoint to test."}