import os
import gradio as gr
from transformers import pipeline
from elasticsearch import Elasticsearch

# Connect to Elasticsearch
es = Elasticsearch(hosts=["https://data.neuralnoise.com:9200"],
                   basic_auth=('elastic', os.environ['ES_PASSWORD']),
                   verify_certs=False, ssl_show_warn=False)

# Load your language model from HuggingFace Transformers
generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")

def search_es(query, index="pubmed", num_results=3):
    """
    Search the Elasticsearch index for the most relevant documents.
    """
    
    print(f'Running query: {query}')
    
    response = es.search(
        index=index,
        body={
            "query": {
                "match": {
                    "content": query  # Assuming documents have a 'content' field
                }
            },
            "size": num_results
        }
    )
    
    # Extract and return the documents
    docs = [hit["_source"]["content"] for hit in response['hits']['hits']]

    print(f'Received {len(docs)} documents')

    return docs

def rag_pipeline(prompt, index="pubmed"):
    """
    A simple RAG pipeline that retrieves documents and uses them to enrich the context for the LLM.
    """
    # Retrieve documents
    docs = search_es(prompt, index=index)
    
    # Combine prompt with retrieved documents
    enriched_prompt = f"{prompt}\n\n{' '.join(docs)}"
    
    # Generate response using the LLM
    response = generator(enriched_prompt, max_new_tokens=256, return_full_text=False)
    
    # Return the generated text and the documents
    return response[0]['generated_text'], "\n\n".join(docs)

# Create the Gradio interface
iface = gr.Interface(fn=rag_pipeline,
                     inputs=[
                         gr.Textbox(label="Input Prompt"),
                         gr.Textbox(label="Elasticsearch Index", value="pubmed")  # Corrected here
                     ],
                     outputs=[
                         gr.Textbox(label="Generated Text"),
                         gr.Textbox(label="Retrieved Documents")
                     ],
                     description="Retrieval-Augmented Generation Pipeline")

# Launch the interface
iface.launch()