Update rag_server.py

rag_server.py

@@ -1,207 +1,121 @@
+import gradio as gr
 import os
-import re
-import glob
-import time
-from collections import defaultdict
-
+import requests
+from io import BytesIO
+from PyPDF2 import PdfReader
+from tempfile import NamedTemporaryFile
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_core.documents import Document
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
-from fastapi.staticfiles import StaticFiles
-from pydantic import BaseModel
-from rag_system import build_rag_chain, ask_question
-from vector_store import get_embeddings, load_vector_store
-from llm_loader import load_llama_model
-import uuid
-from urllib.parse import urljoin, quote
-
-from fastapi.responses import StreamingResponse
-import json
-import time
-
-app = FastAPI()
-
-# Configuration for serving static files
-os.makedirs("static/documents", exist_ok=True)
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-# Prepare global objects
-hf_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-embeddings = get_embeddings(model_name="intfloat/multilingual-e5-large-instruct", device="cpu")
-vectorstore = load_vector_store(embeddings, load_path="vector_db")
-llm = load_llama_model(model_name=hf_model_name) #Hier wird jetzt das Model übergeben
-qa_chain = build_rag_chain(llm, vectorstore, language="en", k=7)
-
-# Server URL configuration (adjust to match your actual environment)
-BASE_URL = "http://220.124.155.35:8500"
-
-class Question(BaseModel):
-    question: str
-
-def get_document_url(source_path):
-    if not source_path or source_path == 'N/A':
-        return None
-    filename = os.path.basename(source_path)
-    dataset_root = os.path.join(os.getcwd(), "dataset")
-    # Find file matching filename in the entire dataset subdirectory
-    found_path = None
-    for root, dirs, files in os.walk(dataset_root):
-        if filename in files:
-            found_path = os.path.join(root, filename)
-            break
-    if not found_path or not os.path.exists(found_path):
-        return None
-    static_path = f"static/documents/{filename}"
-    shutil.copy2(found_path, static_path)
-    encoded_filename = quote(filename)
-    return urljoin(BASE_URL, f"/static/documents/{encoded_filename}")
-
-def create_download_link(url, filename):
-    return f'Source: [{filename}]({url})'
-
-@app.post("/ask")
-def ask(question: Question):
-    result = ask_question(qa_chain, question.question)
-    
-    # Process source document information
-    sources = []
-    for doc in result["source_documents"]:
-        source_path = doc.metadata.get('source', 'N/A')
-        document_url = get_document_url(source_path) if source_path != 'N/A' else None
+from huggingface_hub import InferenceClient
+from gradio.exceptions import Error
+from transformers import AutoModel
+import streamlit as st
+
+# --- Konfiguration ---
+HF_API_TOKEN = os.environ.get("HF_API_TOKEN")  # Lesen Sie den Token aus der Umgebungsvariable
+MODEL_NAME = "dannyk97/mistral-screenplay-model"
+
+# --- Hilfsfunktionen ---
+
+def query_huggingface_inference_endpoints(prompt):
+    """
+    Stellt eine Anfrage an die Hugging Face Inference API.
+    """
+    try:
+        client = InferenceClient(token=HF_API_TOKEN)
+        result = client.text_generation(prompt, model=MODEL_NAME)
+        return result
+    except Exception as e:
+        return f"Fehler bei der Anfrage an Hugging Face API: {e}"
+
+# Function to download PDF from Google Drive
+def download_pdf_from_drive(drive_link):
+    file_id = drive_link.split('/d/')[1].split('/')[0]
+    download_url = f"https://drive.google.com/uc?id={file_id}&export=download"
+    response = requests.get(download_url)
+    if response.status_code == 200:
+        return BytesIO(response.content)
+    else:
+        raise Exception("Failed to download the PDF file from Google Drive.")
+
+# Function to extract text from a PDF
+def extract_text_from_pdf(pdf_stream):
+    pdf_reader = PdfReader(pdf_stream)
+    text = ""
+    for page in pdf_reader.pages:
+        text += page.extract_text()
+    return text
+
+# Function to split text into chunks
+def chunk_text(text, chunk_size=500, chunk_overlap=50):
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        length_function=len
+    )
+    return text_splitter.split_text(text)
+
+# Function to create embeddings and store in FAISS
+def create_embeddings_and_store(chunks):
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    vector_db = FAISS.from_texts(chunks, embedding=embeddings)
+    return vector_db
+
+# Function to query the vector database and interact with Hugging Face Inference API
+def query_vector_db(query, vector_db):
+    # Retrieve relevant documents
+    docs = vector_db.similarity_search(query, k=3)
+    context = "\n".join([doc.page_content for doc in docs])
+
+    # Interact with the Text Generation API
+    prompt = f"Nutze diesen Kontext um die Frage zu beantworten: {context}\nFrage: {query}"
+    try:
+        output = query_huggingface_inference_endpoints(prompt) # Hier wurde das Modell nicht angegeben
+        return output
+    except Exception as e:
+        return f"FEHLER: {str(e)}"
+
+# Streamlit app
+st.title("RAG-Based Application with Google Drive Support")
+
+# Predefined list of Google Drive links - HIER DEFINIERT!
+drive_links = [
+    "https://drive.google.com/file/d/1PW8PJQC1EqYpsk8AhqrE4OS5cy57sqJ4/view?usp=drive_link"
+    # Add more links here as needed
+]
+
+st.write("Processing the predefined Google Drive links...")
+
+all_chunks = []
+
+# Process each predefined Google Drive link
+for link in drive_links:
+    try:
+        st.write(f"Processing link: {link}")
+        # Download PDF
+        pdf_stream = download_pdf_from_drive(link)
+        st.write("PDF Downloaded Successfully!")
         
-        source_info = {
-            "source": source_path,
-            "content": doc.page_content,
-            "page": doc.metadata.get('page', 'N/A'),
-            "document_url": document_url,
-            "filename": os.path.basename(source_path) if source_path != 'N/A' else None
-        }
-        sources.append(source_info)
-    
-    return {
-        "answer": result['result'].split("A:")[-1].strip() if "A:" in result['result'] else result['result'].strip(),
-        "sources": sources
-    }
-
-@app.get("/v1/models")
-def list_models():
-    return JSONResponse({
-        "object": "list",
-        "data": [
-            {
-                "id": "rag",
-                "object": "model",
-                "owned_by": "local",
-            }
-        ]
-    })
-
-@app.post("/v1/chat/completions")
-async def openai_compatible_chat(request: Request):
-    payload = await request.json()
-    messages = payload.get("messages", [])
-    user_input = messages[-1]["content"] if messages else ""
-    stream = payload.get("stream", False)
-
-    result = ask_question(qa_chain, user_input)
-    answer = result['result']
-    
-    # Process source document information
-    sources = []
-    for doc in result["source_documents"]:
-        source_path = doc.metadata.get('source', 'N/A')
-        document_url = get_document_url(source_path) if source_path != 'N/A' else None
-        filename = os.path.basename(source_path) if source_path != 'N/A' else None
+        # Extract text
+        text = extract_text_from_pdf(pdf_stream)
+        st.write("PDF Text Extracted Successfully!")
         
-        source_info = {
-            "source": source_path,
-            "content": doc.page_content,
-            "page": doc.metadata.get('page', 'N/A'),
-            "document_url": document_url,
-            "filename": filename
-        }
-        sources.append(source_info)
-
-    # Output source information one line at a time
-    sources_md = "\nReferences Documents:\n"
-    seen = set()
-    for source in sources:
-        key = (source['filename'], source['document_url'])
-        if source['document_url'] and source['filename'] and key not in seen:
-            sources_md += f"Source: [{source['filename']}]({source['document_url']})\n"
-            seen.add(key)
-
-    final_answer = answer.split("A:")[-1].strip() if "A:" in answer else answer.strip()
-    final_answer += sources_md
-
-    if not stream:
-        return JSONResponse({
-            "id": f"chatcmpl-{uuid.uuid4()}",
-            "object": "chat.completion",
-            "choices": [{
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": final_answer
-                },
-                "finish_reason": "stop"
-            }],
-            "model": "rag",
-        })
-
-    # Generator for streaming response
-    def event_stream():
-        # Stream only the answer body first
-        answer_main = answer.split("A:")[-1].strip() if "A:" in answer else answer.strip()
-        for char in answer_main:
-            chunk = {
-                "id": f"chatcmpl-{uuid.uuid4()}",
-                "object": "chat.completion.chunk",
-                "choices": [{
-                    "index": 0,
-                    "delta": {
-                        "content": char
-                    },
-                    "finish_reason": None
-                }]
-            }
-            yield f"data: {json.dumps(chunk)}\n\n"
-            time.sleep(0.005)
-        # Send reference documents (download links) all at once at the end
-        sources_md = "\nReferences Documents:\n"
-        seen = set()
-        for source in sources:
-            key = (source['filename'], source['document_url'])
-            if source['document_url'] and source['filename'] and key not in seen:
-                sources_md += f"Source: [{source['filename']}]({source['document_url']})\n"
-                seen.add(key)
-        if sources_md.strip() != "References Documents:":
-            chunk = {
-                "id": f"chatcmpl-{uuid.uuid4()}",
-                "object": "chat.completion.chunk",
-                "choices": [{
-                    "index": 0,
-                    "delta": {
-                        "content": sources_md
-                    },
-                    "finish_reason": None
-                }]
-            }
-            yield f"data: {json.dumps(chunk)}\n\n"
-        done = {
-            "id": f"chatcmpl-{uuid.uuid4()}",
-            "object": "chat.completion.chunk",
-            "choices": [{
-                "index": 0,
-                "delta": {},
-                "finish_reason": "stop"
-            }]
-        }
-        yield f"data: {json.dumps(done)}\n\n"
-        return
-
-    return StreamingResponse(event_stream(), media_type="text/event-stream")
+        # Chunk text
+        chunks = chunk_text(text)
+        st.write(f"Created {len(chunks)} text chunks.")
+        all_chunks.extend(chunks)
+    except Exception as e:
+        st.write(f"Error processing link {link}: {e}")
+
+if all_chunks:
+    # Generate embeddings and store in FAISS
+    vector_db = create_embeddings_and_store(all_chunks)
+    st.write("Embeddings Generated and Stored Successfully!")
+    
+    # User query input
+    user_query = st.text_input("Enter your query:")
+    if user_query:
+        response = query_vector_db(user_query, vector_db)
+        st.write("Response from LLM:")
+        st.write(response)