app.py

import os
import shutil
import requests
import json
import gradio as gr
import PyPDF2
import chromadb
import csv
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings

# Constants
API_KEY = os.getenv("togetherai")
BASE_URL = "https://api.together.xyz/v1/chat/completions"
CHUNK_SIZE = 6000  # Maximum words per chunk
TEMP_SUMMARY_FILE = "temp_summaries.txt"
COLLECTIONS_FILE = "collections.csv"

# Function to convert PDF to text
def pdf_to_text(file_path):
    with open(file_path, 'rb') as pdf_file:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text

# Function to summarize text using LLM
def summarize_text(text):
    user_prompt = f"""
    You are an expert in legal language and document summarization. Your task is to provide a concise and accurate summary of the given document. 
    Keep the summary concise, ideally in 2000 words, while covering all essential points. Here is the document to summarize:

    {text}
    """

    return call_llm(user_prompt)

# Function to handle file upload, summarization, and saving to ChromaDB
def handle_file_upload(files, collection_name):
    if not collection_name:
        return "Please provide a collection name."

    os.makedirs('uploaded_pdfs', exist_ok=True)
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=100)
    embeddings = HuggingFaceEmbeddings(model_name="thenlper/gte-small")

    client = chromadb.PersistentClient(path="./db")
    try:
        collection = client.create_collection(name=collection_name)
    except ValueError as e:
        return f"Error creating collection: {str(e)}. Please try a different collection name."

    file_names = []
    with open(TEMP_SUMMARY_FILE, 'w', encoding='utf-8') as temp_file:
        for file in files:
            file_name = os.path.basename(file.name)
            file_names.append(file_name)
            file_path = os.path.join('uploaded_pdfs', file_name)
            shutil.copy(file.name, file_path)
            
            text = pdf_to_text(file_path)
            chunks = text_splitter.split_text(text)
            
            for i, chunk in enumerate(chunks):
                summary = summarize_text(chunk)
                temp_file.write(f"Summary of {file_name} (Part {i+1}):\n{summary}\n\n")

    # Process the temporary file and add to ChromaDB
    with open(TEMP_SUMMARY_FILE, 'r', encoding='utf-8') as temp_file:
        summaries = temp_file.read()
        summary_chunks = text_splitter.split_text(summaries)

        for i, chunk in enumerate(summary_chunks):
            vector = embeddings.embed_query(chunk)
            collection.add(
                embeddings=[vector],
                documents=[chunk],
                ids=[f"summary_{i}"]
            )

    os.remove(TEMP_SUMMARY_FILE)

    # Update collections.csv
    update_collections_csv(collection_name, file_names)

    return "Files uploaded, summarized, and processed successfully."

# Function to update collections.csv
def update_collections_csv(collection_name, file_names):
    file_names_str = ", ".join(file_names)
    with open(COLLECTIONS_FILE, 'a', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow([collection_name, file_names_str])

# Function to read collections.csv
def read_collections():
    if not os.path.exists(COLLECTIONS_FILE):
        return "No collections found."
    
    with open(COLLECTIONS_FILE, 'r') as csvfile:
        reader = csv.reader(csvfile)
        collections = [f"Collection: {row[0]}\nFiles: {row[1]}\n\n" for row in reader]
    
    return "".join(collections)

# Function to search vector database
def search_vector_database(query, collection_name):
    if not collection_name:
        return "Please provide a collection name."

    embeddings = HuggingFaceEmbeddings(model_name="thenlper/gte-small")
    client = chromadb.PersistentClient(path="./db")
    try:
        collection = client.get_collection(name=collection_name)
    except ValueError as e:
        return f"Error accessing collection: {str(e)}. Make sure the collection name is correct."

    query_vector = embeddings.embed_query(query)
    results = collection.query(query_embeddings=[query_vector], n_results=2, include=["documents"])
    
    return "\n\n".join(results["documents"][0])

# Function to call LLM
def call_llm(prompt):
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    data = {
        "model": "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.7,
        "top_p": 0.7,
        "top_k": 50,
        "repetition_penalty": 1,
        "stop": ["\"\""],
        "stream": False
    }

    response = requests.post(BASE_URL, headers=headers, data=json.dumps(data))
    response.raise_for_status()
    return response.json()['choices'][0]['message']['content']

# Function to answer questions using Rachel.AI
def answer_question(question, collection_name):
    context = search_vector_database(question, collection_name)
    
    prompt = f"""
    You are a paralegal AI assistant. Your role is to assist with legal inquiries by providing clear and concise answers based on the provided question and legal context. Always maintain a highly professional tone, ensuring that your responses are well-reasoned and legally accurate.
    Question: {question}
    Legal Context: {context}
    Please provide a detailed response considering the above information. also when answering the question make sure to inform the user which filename you are using (filename is often given in the context)
    """
    
    return call_llm(prompt)

# Gradio interface
def gradio_interface():
    with gr.Blocks(theme='gl198976/The-Rounded') as interface:
        gr.Markdown("# rachel.ai backend")

        gr.Markdown("""
        ### Warning
        If you encounter an error when uploading files, try changing the collection name and upload again.
        Each collection name must be unique.
        """)
        
        with gr.Tab("Document Upload and Search"):
            with gr.Row():
                with gr.Column():
                    collection_name_input = gr.Textbox(label="Collection Name", placeholder="Enter a unique name for this collection")
                    file_upload = gr.Files(file_types=[".pdf"], label="Upload PDFs")
                    upload_btn = gr.Button("Upload, Summarize, and Process Files")
                    upload_status = gr.Textbox(label="Upload Status", interactive=False)
                with gr.Column():
                    search_query_input = gr.Textbox(label="Search Query")
                    search_collection_name = gr.Textbox(label="Collection Name for Search", placeholder="Enter the collection name to search")
                    search_output = gr.Textbox(label="Search Results", lines=10)
                    search_btn = gr.Button("Search")
            
            api_details = gr.Markdown("""
                ### API Endpoint Details
                - **URL:** http://0.0.0.0:7860/search_vector_database
                - **Method:** POST
                - **Example Usage:**
                
                ```python
                from gradio_client import Client

                client = Client("http://0.0.0.0:7860/")
                result = client.predict(
                    "search query",  # str in 'Search Query' Textbox component
                    "name of collection given in ui",  # str in 'Collection Name' Textbox component
                    api_name="/search_vector_database"
                )
                print(result)
                ```
            """)

        with gr.Tab("Rachel.AI"):
            question_input = gr.Textbox(label="Ask a question")
            rachel_collection_name = gr.Textbox(label="Collection Name", placeholder="Enter the collection name to search")
            answer_output = gr.Textbox(label="Answer", lines=10)
            ask_btn = gr.Button("Ask Rachel.AI")
            
            rachel_api_details = gr.Markdown("""
                ### API Endpoint Details for Rachel.AI
                - **URL:** http://0.0.0.0:7860/answer_question
                - **Method:** POST
                - **Example Usage:**
                
                ```python
                from gradio_client import Client

                client = Client("http://0.0.0.0:7860/")
                result = client.predict(
                    "question",  # str in 'Ask a question' Textbox component
                    "collection_name",  # str in 'Collection Name' Textbox component
                    api_name="/answer_question"
                )
                print(result)
                ```
            """)

        with gr.Tab("Collections"):
            collections_output = gr.Textbox(label="Collections and Files", lines=20)
            refresh_btn = gr.Button("Refresh Collections")

        upload_btn.click(handle_file_upload, inputs=[file_upload, collection_name_input], outputs=[upload_status])
        search_btn.click(search_vector_database, inputs=[search_query_input, search_collection_name], outputs=[search_output])
        ask_btn.click(answer_question, inputs=[question_input, rachel_collection_name], outputs=[answer_output])
        refresh_btn.click(read_collections, inputs=[], outputs=[collections_output])

    interface.launch(server_name="0.0.0.0", server_port=7860)

if __name__ == "__main__":
    gradio_interface()