app.py

# === Required Libraries ===
from huggingface_hub import InferenceClient
import os
import sys
import re
import json
import requests
import logging
from bs4 import BeautifulSoup
from readability import Document
from duckduckgo_search import DDGS
from concurrent.futures import ThreadPoolExecutor
import gradio as gr
from datetime import datetime, timedelta
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import wikipedia

# === Configuration ===
HF_TOKEN = os.getenv("HF") 
MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
client = InferenceClient(model=MODEL_NAME, token=HF_TOKEN)

HEADERS = {"User-Agent": "Mozilla/5.0"}
MAX_RESULTS = 5
MAX_CHARS = 5000
CONTEXT_DIR = "web_contexts"
CHUNK_STORE = "chunk_store.json"
LOG_FILE = "qa_log.jsonl"
EMBED_FILE = "memory_embeddings.json"
MAX_CHUNK_AGE_DAYS = 3
MIN_CONTEXT_SIMILARITY = 0.4

SEMANTIC_SCHOLAR_API = "https://api.semanticscholar.org/graph/v1/paper/search"
SEMANTIC_SCHOLAR_FIELDS = "title,abstract,url,authors,year"

os.makedirs(CONTEXT_DIR, exist_ok=True)
logging.basicConfig(level=logging.INFO)

# === Embedding Model ===
EMBED_MODEL = SentenceTransformer("all-MiniLM-L6-v2")

def embed(text):
    emb = EMBED_MODEL.encode([text])[0]
    return np.array(emb, dtype=np.float32)

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-10)

def chunk_text(text, max_tokens=200):
    sentences = re.split(r'(?<=[.!?]) +', text)
    chunks, chunk, tokens = [], [], 0
    for sent in sentences:
        sent_tokens = len(sent.split())
        if tokens + sent_tokens > max_tokens:
            chunks.append(" ".join(chunk))
            chunk, tokens = [], 0
        chunk.append(sent)
        tokens += sent_tokens
    if chunk:
        chunks.append(" ".join(chunk))
    return chunks

def current_iso_timestamp():
    return datetime.utcnow().isoformat()

def save_chunks(query, chunks, urls):
    chunk_data = []
    now = current_iso_timestamp()
    for chunk in chunks:
        chunk_data.append({
            "query": query,
            "chunk": chunk,
            "embedding": embed(chunk).tolist(),
            "sources": urls,
            "timestamp": now
        })
    if os.path.exists(CHUNK_STORE):
        with open(CHUNK_STORE, "r") as f:
            existing = json.load(f)
        cutoff = datetime.utcnow() - timedelta(days=MAX_CHUNK_AGE_DAYS)
        existing = [c for c in existing if datetime.fromisoformat(c.get("timestamp", "1970-01-01T00:00:00")) > cutoff]
    else:
        existing = []
    existing.extend(chunk_data)
    with open(CHUNK_STORE, "w") as f:
        json.dump(existing, f, indent=2)

def is_recent_chunk(ts):
    try:
        return (datetime.utcnow() - datetime.fromisoformat(ts)) < timedelta(days=MAX_CHUNK_AGE_DAYS)
    except:
        return False

def retrieve_context_from_chunks(question, top_k=4):
    if not os.path.exists(CHUNK_STORE):
        return "", [], 0.0
    with open(CHUNK_STORE, "r") as f:
        data = json.load(f)
    data = [d for d in data if is_recent_chunk(d.get("timestamp", ""))]
    if not data:
        return "", [], 0.0

    embeddings = np.array([d['embedding'] for d in data]).astype('float32')
    dim = embeddings.shape[1]
    q_emb = embed(question).reshape(1, -1).astype('float32')
    if q_emb.shape[1] != dim:
        os.remove(CHUNK_STORE)
        return "", [], 0.0

    index = faiss.IndexFlatL2(dim)
    index.add(embeddings)
    distances, I = index.search(q_emb, top_k)
    top_chunks = [data[i]['chunk'] for i in I[0]]
    sources = list({src for i in I[0] for src in data[i]['sources']})
    similarities = 1 / (distances[0] + 1e-6)
    avg_sim = np.mean(similarities)
    return "\n\n".join(top_chunks), sources, avg_sim

def fetch_text(url):
    try:
        r = requests.get(url, headers=HEADERS, timeout=10)
        doc = Document(r.text)
        soup = BeautifulSoup(doc.summary(), "html.parser")
        text = " ".join(p.get_text() for p in soup.find_all("p")).strip()
        return text, url
    except Exception as e:
        return "", url

def scrape_and_save(query):
    filename = re.sub(r'[^a-zA-Z0-9_-]', '_', query)[:50] + ".json"
    filepath = os.path.join(CONTEXT_DIR, filename)
    if os.path.exists(filepath):
        with open(filepath, "r") as f:
            d = json.load(f)
        return d["context"], d["sources"]

    with DDGS() as ddgs:
        results = list(ddgs.text(query, max_results=MAX_RESULTS))

    urls = list({r['href'] for r in results if 'href' in r})
    with ThreadPoolExecutor(max_workers=MAX_RESULTS) as executor:
        fetched = list(executor.map(fetch_text, urls))

    texts, used_urls, total_chars = [], [], 0
    q_emb = embed(query)
    for text, url in fetched:
        if not text:
            continue
        if query.lower() not in text.lower():
            sim = cosine_similarity(q_emb, embed(text))
            if sim < 0.3:
                continue
        if total_chars + len(text) > MAX_CHARS:
            text = text[:MAX_CHARS - total_chars]
        texts.append(text)
        used_urls.append(url)
        total_chars += len(text)
        if total_chars >= MAX_CHARS:
            break

    context = "\n\n".join(texts)
    chunks = chunk_text(context)
    save_chunks(query, chunks, used_urls)
    with open(filepath, "w") as f:
        json.dump({"query": query, "context": context, "sources": used_urls}, f, indent=2)
    return context, used_urls

def get_similar_memories(question, top_k=3):
    if not os.path.exists(EMBED_FILE):
        return []
    with open(EMBED_FILE, "r") as f:
        data = json.load(f)
    if not data:
        return []

    embeddings = np.array([m['embedding'] for m in data]).astype('float32')
    dim = embeddings.shape[1]
    q_emb = embed(question).reshape(1, -1).astype('float32')
    if q_emb.shape[1] != dim:
        os.remove(EMBED_FILE)
        return []

    index = faiss.IndexFlatL2(dim)
    index.add(embeddings)
    _, I = index.search(q_emb, top_k)
    return [data[i] for i in I[0]]

def save_embedding_to_store(entry):
    if os.path.exists(EMBED_FILE):
        with open(EMBED_FILE, "r") as f:
            data = json.load(f)
    else:
        data = []
    data.append(entry)
    with open(EMBED_FILE, "w") as f:
        json.dump(data, f, indent=2)

def answer_from_context(question):
    memory = get_similar_memories(question)
    memory_prompt = "\n\n".join(f"Q: {m['q']}\nA: {m['a']}" for m in memory)
    context, sources, avg_sim = retrieve_context_from_chunks(question)

    prompt = f"""
Today's date is {datetime.utcnow().date()}.
Use context and memory to answer and summarize the following question using fullly finished lines end with., clear, and grammatically correct finish sentences. Ensure that the response is factually accurate, complete, well-organized, finish sentences and easy to understand. Avoid repeating information,unfinish sentences and keep the response concise while still being informative.

[CONTEXT]
{context}

[MEMORY]
{memory_prompt}

[QUESTION]
Answer and summarize the following question using fullly finish linesens end with., clear, and grammatically correct finish sentences. Ensure that the response is factually accurate, complete, well-organized, finish sentences and easy to understand. Avoid repeating information, unfinish sentences, and keep the response concise while still being informative.
{question}

[ANSWER]
"""
    try:
        response = client.text_generation(prompt, max_new_tokens=512)
        reply = response.strip().split("<|assistant|>")[-1].strip()
    except Exception as e:
        reply = f"Error: {e}"

    log = {
        "time": str(datetime.utcnow()),
        "q": question,
        "a": reply,
        "sources": sources,
        "embedding": embed(question).tolist()
    }
    with open(LOG_FILE, "a") as f:
        f.write(json.dumps(log) + "\n")
    save_embedding_to_store(log)
    return reply, sources, avg_sim

def needs_web_search_llm(question):
    prompt = f"""
You are a helpful assistant that classifies whether a question requires a web search or external data.

Question: "{question}"

Answer with only "YES" if a web search is needed or "NO" if not.
"""
    try:
        response = client.text_generation(prompt, max_new_tokens=10)
        return "YES" in response.strip().upper()
    except Exception as e:
        return False

def is_general_knowledge_question(question):
    prompt = f"""
You are a classifier. Determine if the question below can be answered using general world knowledge, like an encyclopedia or Wikipedia.

Question: "{question}"

Answer with "YES" if it is general knowledge. Otherwise answer "NO".
"""
    try:
        response = client.text_generation(prompt, max_new_tokens=10)
        return "YES" in response.strip().upper()
    except Exception as e:
        return False

def get_wikipedia_summary(query, sentences=3):
    try:
        wikipedia.set_lang("en")
        return wikipedia.summary(query, sentences=sentences)
    except wikipedia.exceptions.DisambiguationError as e:
        return f"Ambiguous question. Possible topics: {', '.join(e.options[:5])}"
    except wikipedia.exceptions.PageError:
        return "No Wikipedia article found for that topic."
    except Exception as e:
        return "Error accessing Wikipedia."

# === Semantic Scholar API integration ===
def semantic_scholar_search(query, max_results=5):
    params = {
        "query": query,
        "fields": SEMANTIC_SCHOLAR_FIELDS,
        "limit": max_results
    }
    try:
        resp = requests.get(SEMANTIC_SCHOLAR_API, params=params, timeout=10)
        resp.raise_for_status()
        data = resp.json()
        papers = data.get("data", [])
        texts = []
        urls = []
        for p in papers:
            title = p.get("title", "")
            abstract = p.get("abstract", "")
            url = p.get("url", "")
            year = p.get("year", "")
            authors = ", ".join([a.get("name","") for a in p.get("authors", [])])
            entry = f"Title: {title}\nAuthors: {authors}\nYear: {year}\nAbstract: {abstract}\nURL: {url}\n"
            texts.append(entry)
            if url:
                urls.append(url)
            if len("\n\n".join(texts)) > MAX_CHARS:
                break
        context = "\n\n".join(texts)
        chunks = chunk_text(context)
        save_chunks(query, chunks, urls)
        return context, urls
    except Exception as e:
        logging.warning(f"Semantic Scholar API error: {e}")
        return "", []

def is_research_question(question):
    # Simple heuristic to detect research/scientific questions
    keywords = [
    "research", "study", "paper", "findings", "experiment", "scientific", "evidence", "meta-analysis",
    "hypothesis", "literature review", "case study", "theory", "framework", "methodology", "analysis",
    "data", "observation", "results", "variables", "survey", "questionnaire", "sampling", "experiment design",
    "quantitative", "qualitative", "mixed methods", "statistical", "inference", "regression", "correlation",
    "interview", "focus group", "coding", "themes", "interpretation", "reliability", "validity", "bias",
    "significance", "conclusion", "discussion", "implications", "limitations", "future research", "peer review",
    "publication", "citation", "replication", "protocol", "ethics", "IRB", "research question", "objective",
    "aim", "problem statement", "gap", "contribution", "novelty", "originality", "dataset", "case", "fieldwork",
    "observational", "experimental", "review", "systematic review", "control group", "randomized", "longitudinal",
    "cross-sectional", "data analysis", "research design", "conceptual", "empirical", "exploratory", "descriptive",
    "causal", "predictive", "construct", "operationalization", "dependent variable", "independent variable",
    "mediator", "moderator", "association", "impact", "effect", "relationship", "outcome", "measure", "coding scheme"
]

    q_lower = question.lower()
    return any(kw in q_lower for kw in keywords)

def ask(q):
    # Check if research/scientific question and use Semantic Scholar
    if is_research_question(q):
        context, sources = semantic_scholar_search(q)
        if context:
            answer, sources, _ = answer_from_context(q)
            sources_text = "\n".join(f"- {url}" for url in sources)
            return answer, sources_text
        # fallback to regular web search if semantic scholar fails
        context, sources = scrape_and_save(q)
        answer, sources, _ = answer_from_context(q)
        sources_text = "\n".join(f"- {url}" for url in sources)
        return answer, sources_text

    # General knowledge questions use Wikipedia
    if is_general_knowledge_question(q):
        return get_wikipedia_summary(q), "Source: Wikipedia"

    # Check if we already have context stored with sufficient similarity
    _, _, avg_sim = retrieve_context_from_chunks(q)

    # Check if web search is needed or context similarity too low
    intent_search = needs_web_search_llm(q)

    if intent_search or avg_sim < MIN_CONTEXT_SIMILARITY:
        context, sources = scrape_and_save(q)
        answer, sources, _ = answer_from_context(q)
        sources_text = "\n".join(f"- {url}" for url in sources)
    else:
        # Use model to answer from prompt only
        prompt = f"<|user|>\n Answer  and summarize the following question using fullly finish lines end with. , clear, and grammatically correct finish sentences. Ensure that the response is factually accurate, complete, well-organized, finish stances, and easy to understand. Avoid repeating information, unfinish sentences, and keep the response concise while still being informative.:\n{q.strip()}\n<|assistant|>\n"
        try:
            response = client.text_generation(prompt, max_new_tokens=512)
            answer = response.strip().split("<|assistant|>")[-1].strip()
        except Exception as e:
            answer = f"Error: {e}"
        sources_text = ""

    return answer, sources_text

# === Gradio UI ===
with gr.Blocks() as demo:
    gr.Markdown("""
## 🤖 LLaMA 3.1 Smart QA Bot  
- Uses **Wikipedia** for general knowledge  
- Searches **Semantic Scholar** for research-related questions  
- Falls back to web search when needed  
- Can handle **casual chat** too!
""")

    q_input = gr.Textbox(label="Your Question")
    submit = gr.Button("Ask")
    a_output = gr.Textbox(label="Answer")
    s_output = gr.Markdown()
    submit.click(ask, inputs=q_input, outputs=[a_output, s_output])

if __name__ == '__main__':
    if len(sys.argv) > 1:
        question = " ".join(sys.argv[1:])
        print(ask(question))
    else:
        demo.launch()