app.py

import os
from dotenv import load_dotenv
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import HumanMessage
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_voyageai import VoyageAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain.prompts import PromptTemplate
from pinecone import Pinecone
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import openai
import gradio as gr

# Load API keys
load_dotenv()
openai.api_key = os.environ.get("OPENAI_API_KEY")
pinecone_api_key = os.environ.get("PINECONE_API_KEY")
voyage_api_key = os.environ.get("VOYAGE_API_KEY")

# Initialize Pinecone
pc = Pinecone(api_key=pinecone_api_key)
embeddings = VoyageAIEmbeddings(voyage_api_key=voyage_api_key, model="voyage-law-2")

# 🔹 Query Expansion using GPT-4
def expand_query(query):
    llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.3)
    prompt = f"Rewrite this vague query for searching a document into a more specific one:\nQuery: {query}\nSpecific Query:"
    refined_query = llm([HumanMessage(content=prompt)]).content.strip()
    return refined_query if refined_query else query

# 🔹 Hybrid Search (TF-IDF + Semantic Retrieval)
def hybrid_search(query, user_groups, index_name="briefmeta", min_score=0.01, fetch_k=50):
    vector_store = PineconeVectorStore(index_name=index_name, embedding=embeddings)
    semantic_results = vector_store.max_marginal_relevance_search(query, k=10, fetch_k=fetch_k)

    all_texts = [doc.page_content for doc in semantic_results]
    vectorizer = TfidfVectorizer(stop_words="english")
    tfidf_matrix = vectorizer.fit_transform(all_texts)
    query_tfidf = vectorizer.transform([query])
    keyword_scores = cosine_similarity(query_tfidf, tfidf_matrix).flatten()

    combined_results, seen_ids = [], set()
    for i, doc in enumerate(semantic_results):
        doc_id, doc_groups = doc.metadata.get("id"), doc.metadata.get("groups", [])
        semantic_score = float(doc.metadata.get("score", 0))
        keyword_score = float(keyword_scores[i])
        final_score = 0.65 * semantic_score + 0.35 * keyword_score  # Hybrid score

        if doc_id not in seen_ids and any(group in user_groups for group in doc_groups) and final_score > min_score:
            seen_ids.add(doc_id)
            doc.metadata["final_score"] = final_score
            combined_results.append(doc)

    combined_results.sort(key=lambda x: x.metadata["final_score"], reverse=True)
    return [
        {
            "doc_id": doc.metadata.get("doc_id", "N/A"),
            "chunk_id": doc.metadata.get("id", "N/A"),
            "title": doc.metadata.get("source", "N/A"),
            "text": doc.page_content,
            "page_number": str(doc.metadata.get("page_number", "N/A")),
            "score": str(doc.metadata.get("final_score", "N/A")),
        }
        for doc in combined_results
    ]

# 🔹 Metadata-Weighted Reranking
def rerank(query, context):
    reranker = pc.inference.rerank(
        model="bge-reranker-v2-m3", query=query, documents=context, top_n=10, return_documents=True
    )

    final_reranked = []
    for entry in reranker.data:
        doc, score = entry["document"], float(entry["score"])
        citation_boost = 1.2 if "high_citations" in doc.get("tags", []) else 1.0
        recency_boost = 1.1 if "recent_upload" in doc.get("tags", []) else 1.0
        final_score = score * citation_boost * recency_boost
        doc["final_score"] = final_score
        final_reranked.append(doc)

    final_reranked.sort(key=lambda x: x["final_score"], reverse=True)
    return final_reranked

# 🔹 Intelligent Search Summary Generator
def generate_search_summary(search_results, document_titles, query):
    if not search_results:
        return "No relevant documents found. Try refining your query."

    num_results = len(document_titles)
    doc_titles = [doc.get("title", "Unknown Document") for doc in search_results]
    doc_pages = [doc.get("page_number", "N/A") for doc in search_results]
    relevance_scores = [float(doc.get("score", 0)) for doc in search_results]

    summary_prompt = f"""
    Generate a concise 1-3 sentence summary for the document search results found:
    - User Query: "{query}"
    - Matching Documents: {num_results} found
    - Titles: {", ".join(set(doc_titles))}
    - Pages Referenced: {", ".join(set(doc_pages))}
    - Relevance Scores (0-1): {relevance_scores}
    Provide a clear, user-friendly summary with an action suggestion.
    If scores are low but the documents are from the same title no need to comment on the scores.
    If only one result that means there's is only one relevant document and in that case do not mention about pages/page number.
    """

    llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.5)
    summary = llm([HumanMessage(content=summary_prompt)]).content.strip()
    return summary if summary else "No intelligent summary available."

# 🔹 LLM-based Answer Generation
def generate_output(context, query):
    if not context.strip():
        return "No relevant information found. Try refining your query."

    llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.5)
    prompt_template = PromptTemplate(
        template="Use the following context to answer the question:\nContext: {context}\nQuestion: {question}\nAnswer:",
        input_variables=["context", "question"],
    )
    prompt = prompt_template.format(context=context, question=query)
    response = llm([HumanMessage(content=prompt)]).content.strip()
    return response if response else "No relevant answer found."

# 🔹 Full Workflow
def complete_workflow(query, user_groups, index_name="briefmeta"):
    try:
        refined_query = expand_query(query)
        context_data = hybrid_search(refined_query, user_groups)
        reranked_results = rerank(refined_query, context_data)

        context_data = [
            {
                'chunk_id': doc["chunk_id"],
                'doc_id': doc["doc_id"],
                'title': doc["title"],
                'text': doc["text"],
                'page_number': str(doc["page_number"]),
                'score': str(doc["final_score"])
            }
            for doc in reranked_results
        ]

        document_titles = list({os.path.basename(doc["title"]) for doc in context_data})
        formatted_titles = " " + "\n".join(document_titles)
        intelligent_search_summary = generate_search_summary(context_data, document_titles, refined_query)

        results = {
            "results": [
                {
                    "natural_language_output": generate_output(doc["text"], refined_query),
                    "chunk_id": doc["chunk_id"],
                    "document_id": doc["doc_id"],
                    "title": doc["title"],
                    "text": doc["text"],
                    "page_number": doc["page_number"],
                    "score": doc["score"],
                }
                for doc in context_data
            ],
            "total_results": len(context_data),
            "intelligent_search_summary": intelligent_search_summary
        }

        return results, formatted_titles, intelligent_search_summary

    except Exception as e:
        return {"results": [], "total_results": 0, "intelligent_search_summary": "Error generating summary."}, f"Error in workflow: {str(e)}"

# 🔹 Gradio UI
def gradio_app():
    with gr.Blocks() as app:
        gr.Markdown("### 📄 Intelligent Document Search Prototype-v0.2")
        user_query = gr.Textbox(label="🔍 Enter Search Query")
        user_groups = gr.Textbox(label="👥 User Groups", placeholder="e.g., ['KarthikPersonal']")
        index_name = gr.Textbox(label="📂 Index Name", placeholder="Default: briefmeta")
        search_btn = gr.Button("🔎 Search")
        search_summary = gr.Textbox(label="📜 Intelligent Search Summary", interactive=False)
        result_output = gr.JSON(label="📊 Search Results")
        titles_output = gr.Textbox(label="📂 Retrieved Document Titles", interactive=False)

        search_btn.click(complete_workflow, inputs=[user_query, user_groups, index_name], outputs=[result_output, titles_output, search_summary])

    return app

# Launch the App
gradio_app().launch()