app.py

# Combined Llama 3 and Gemini Flash Chatbot
import json
import logging
import re
import os
import pickle
from typing import List, Tuple, Optional
import gradio as gr
from openai import OpenAI
import google.generativeai as genai
from functools import lru_cache
from tenacity import retry, stop_after_attempt, wait_exponential
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import FAISS
from langchain_core.embeddings import Embeddings
from langchain_core.documents import Document
from collections import defaultdict
import hashlib
from tqdm import tqdm
from dotenv import load_dotenv

load_dotenv()

# --- Configuration ---
FAISS_INDEX_PATH = "faiss_index"
BM25_INDEX_PATH = "bm25_index.pkl"
CACHE_VERSION = "v1"
embedding_model = "e5-mistral-7b-instruct"
data_file_name = "AskNatureNet_data_enhanced.json"
CHUNK_SIZE = 800
OVERLAP = 200
EMBEDDING_BATCH_SIZE = 32

# Initialize clients
OPENAI_API_CONFIG = {
    "api_key": os.getenv("OPENAI_API_KEY"),
    "base_url": "https://chat-ai.academiccloud.de/v1"
}
client = OpenAI(**OPENAI_API_CONFIG)
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Helper Functions ---
def get_data_hash(file_path: str) -> str:
    """Generate hash of data file for cache validation"""
    with open(file_path, "rb") as f:
        return hashlib.md5(f.read()).hexdigest()

# --- Custom Embedding Handler ---
class MistralEmbeddings(Embeddings):
    """E5-Mistral-7B embedding adapter"""
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        embeddings = []
        try:
            for i in tqdm(range(0, len(texts), EMBEDDING_BATCH_SIZE), desc="Embedding Progress"):
                batch = texts[i:i + EMBEDDING_BATCH_SIZE]
                response = client.embeddings.create(
                    input=batch,
                    model=embedding_model,
                    encoding_format="float"
                )
                embeddings.extend([e.embedding for e in response.data])
            return embeddings
        except Exception as e:
            logger.error(f"Embedding Error: {str(e)}")
            return [[] for _ in texts]
        
    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]

# --- Data Processing ---
def load_and_chunk_data(file_path: str) -> List[Document]:
    """Enhanced chunking with metadata preservation"""
    current_hash = get_data_hash(file_path)
    cache_file = f"documents_{CACHE_VERSION}_{current_hash}.pkl"
    
    if os.path.exists(cache_file):
        logger.info("Loading cached documents")
        with open(cache_file, "rb") as f:
            return pickle.load(f)
    
    with open(file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    documents = []
    for item in tqdm(data, desc="Chunking Progress"):
        base_content = f"""Source: {item['Source']}
Application: {item['Application']}
Functions: {', '.join(filter(None, [item.get('Function1'), item.get('Function2')]))}
Technical Concepts: {', '.join(item['technical_concepts'])}
Biological Mechanisms: {', '.join(item['biological_mechanisms'])}"""
        
        strategy = item['Strategy']
        for i in range(0, len(strategy), CHUNK_SIZE - OVERLAP):
            chunk = strategy[i:i + CHUNK_SIZE]
            documents.append(Document(
                page_content=f"{base_content}\nStrategy Excerpt:\n{chunk}",
                metadata={
                    "source": item["Source"],
                    "application": item["Application"],
                    "technical_concepts": item["technical_concepts"],
                    "sustainability_impacts": item["sustainability_impacts"],
                    "hyperlink": item["Hyperlink"],
                    "chunk_id": f"{item['Source']}-{len(documents)+1}"
                }
            ))
    
    with open(cache_file, "wb") as f:
        pickle.dump(documents, f)
    return documents

# --- Optimized Retrieval System ---
class EnhancedRetriever:
    """Hybrid retriever with persistent caching"""
    def __init__(self, documents: List[Document]):
        self.documents = documents
        self.bm25 = self._init_bm25()
        self.vector_store = self._init_faiss()
        self.vector_retriever = self.vector_store.as_retriever(search_kwargs={"k": 3})

    def _init_bm25(self) -> BM25Retriever:
        cache_key = f"{BM25_INDEX_PATH}_{get_data_hash(data_file_name)}"
        if os.path.exists(cache_key):
            logger.info("Loading cached BM25 index")
            with open(cache_key, "rb") as f:
                return pickle.load(f)
        
        logger.info("Building new BM25 index")
        retriever = BM25Retriever.from_documents(self.documents)
        retriever.k = 5
        with open(cache_key, "wb") as f:
            pickle.dump(retriever, f)
        return retriever

    def _init_faiss(self) -> FAISS:
        cache_key = f"{FAISS_INDEX_PATH}_{get_data_hash(data_file_name)}"
        if os.path.exists(cache_key):
            logger.info("Loading cached FAISS index")
            return FAISS.load_local(
                cache_key,
                MistralEmbeddings(),
                allow_dangerous_deserialization=True
            )
        
        logger.info("Building new FAISS index")
        vector_store = FAISS.from_documents(self.documents, MistralEmbeddings())
        vector_store.save_local(cache_key)
        return vector_store

    @lru_cache(maxsize=500)
    def retrieve(self, query: str) -> str:
        try:
            processed_query = self._preprocess_query(query)
            expanded_query = self._hyde_expansion(processed_query)
            
            bm25_results = self.bm25.invoke(processed_query)
            vector_results = self.vector_retriever.invoke(processed_query)
            expanded_results = self.bm25.invoke(expanded_query)
            
            fused_results = self._fuse_results([bm25_results, vector_results, expanded_results])
            return self._format_context(fused_results[:5])
        except Exception as e:
            logger.error(f"Retrieval Error: {str(e)}")
            return ""

    def _preprocess_query(self, query: str) -> str:
        return query.lower().strip()

    @lru_cache(maxsize=500)
    def _hyde_expansion(self, query: str) -> str:
        try:
            response = client.chat.completions.create(
                model="meta-llama-3-70b-instruct",
                messages=[{
                    "role": "user",
                    "content": f"Generate a technical draft about biomimicry for: {query}\nInclude domain-specific terms."
                }],
                temperature=0.5,
                max_tokens=200
            )
            return response.choices[0].message.content
        except Exception as e:
            logger.error(f"HyDE Error: {str(e)}")
            return query

    def _fuse_results(self, result_sets: List[List[Document]]) -> List[Document]:
        fused_scores = defaultdict(float)
        for docs in result_sets:
            for rank, doc in enumerate(docs, 1):
                fused_scores[doc.metadata["chunk_id"]] += 1 / (rank + 60)
        
        seen = set()
        return [
            doc for doc in sorted(
                (doc for docs in result_sets for doc in docs),
                key=lambda x: fused_scores[x.metadata["chunk_id"]],
                reverse=True
            ) if not (doc.metadata["chunk_id"] in seen or seen.add(doc.metadata["chunk_id"]))
        ]

    def _format_context(self, docs: List[Document]) -> str:
        context = []
        for doc in docs:
            context_str = f"""**Source**: [{doc.metadata['source']}]({doc.metadata['hyperlink']})
    **Application**: {doc.metadata['application']}
    **Key Concepts**: {', '.join(doc.metadata['technical_concepts'])}
    **Strategy Excerpt**:\n{doc.page_content.split('Strategy Excerpt:')[-1].strip()}"""
            context.append(context_str)
        return "\n\n---\n\n".join(context)

# --- Generation System ---
SYSTEM_PROMPT = """**Biomimicry Expert Guidelines**
1. Base answers strictly on context
2. Cite sources as [Source]
3. **Bold** technical terms
4. Include reference links

Context: {context}"""

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=20))
def get_ai_response(query: str, context: str, model: str) -> str:
    try:
        if model == "gemini-2.0-flash":
            gemini_model = genai.GenerativeModel(model)
            response = gemini_model.generate_content(
                f"{SYSTEM_PROMPT.format(context=context)}\nQuestion: {query}\nProvide a detailed technical answer:"
            )
            return _postprocess_response(response.text)
        elif model == "meta-llama-3-70b-instruct":
            response = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT.format(context=context)},
                    {"role": "user", "content": f"Question: {query}\nProvide a detailed technical answer:"}
                ],
                temperature=0.4,
                max_tokens=2000
            )
            return _postprocess_response(response.choices[0].message.content)
    except Exception as e:
        logger.error(f"Generation Error: {str(e)}")
        return "I'm unable to generate a response right now. Please try again later."

def _postprocess_response(response: str) -> str:
    response = re.sub(r"\[(.*?)\]", r"[\1](#)", response)
    response = re.sub(r"\*\*([\w-]+)\*\*", r"**\1**", response)
    return response

# --- Pipeline ---
documents = load_and_chunk_data(data_file_name)
retriever = EnhancedRetriever(documents)

def generate_response(question: str, model: str) -> str:
    try:
        context = retriever.retrieve(question)
        return get_ai_response(question, context, model) if context else "No relevant information found."
    except Exception as e:
        logger.error(f"Pipeline Error: {str(e)}")
        return "An error occurred processing your request."

# --- Gradio Interface ---
def chat_interface(question: str, history: List[Tuple[str, str]], model: str):
    response = generate_response(question, model)
    return "", history + [(question, response)]

with gr.Blocks(title="AskNature BioRAG Expert", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌿 AskNature RAG-based Chatbot ")
    with gr.Row():
        chatbot = gr.Chatbot(label="Dialogue History", height=500)
    with gr.Row():
        question = gr.Textbox(placeholder="Ask about biomimicry (e.g. 'How does Werewool use coral proteins to make fibers?')",
                            label="Inquiry", scale=4)
        model_selector = gr.Dropdown(choices=["gemini-2.0-flash", "meta-llama-3-70b-instruct"], label="Generation Model", value="gemini-2.0-flash")
        clear_btn = gr.Button("Clear History", variant="secondary")
    
    gr.Markdown("""
        <div style="text-align: center; color: #4a7c59;">
            <small>Powered by AskNature's Database | 
            Explore nature's blueprints at <a href="https://asknature.org">asknature.org</a></small>
        </div>""")
    question.submit(chat_interface, [question, chatbot, model_selector], [question, chatbot])
    clear_btn.click(lambda: [], None, chatbot)

if __name__ == "__main__":
    demo.launch(show_error=True)