Update app.py

app.py

@@ -2,15 +2,19 @@ import gradio as gr
 import chromadb
 from chromadb.utils import embedding_functions
 from PyPDF2 import PdfReader
+from gradio_client import Client
 
 # Initialisiere ChromaDB
-client = chromadb.Client()
+client_chroma = chromadb.Client()
 collection_name = "pdf_collection"
-collection = client.get_or_create_collection(name=collection_name)
+collection = client_chroma.get_or_create_collection(name=collection_name)
 
 # Verwende die integrierten Embeddings von ChromaDB
 embedding_function = embedding_functions.DefaultEmbeddingFunction()
 
+# Initialisiere das LLM
+client_llm = Client("Mxytyu/meta-llama-Llama-3.2-3B")
+
 def process_pdf(file):
     # Lese den PDF-Inhalt
     pdf_reader = PdfReader(file.name)
@@ -30,14 +34,54 @@ def process_pdf(file):
     
     return f"PDF {file.name} wurde erfolgreich in ChromaDB gespeichert."
 
+def search_similar_documents(prompt):
+    # Erstelle Embedding für den Prompt
+    query_embedding = embedding_function([prompt])[0]
+    
+    # Führe die Ähnlichkeitssuche durch
+    results = collection.query(
+        query_embeddings=[query_embedding],
+        n_results=3
+    )
+    
+    # Formatiere die Ergebnisse
+    formatted_results = []
+    for i, doc in enumerate(results["documents"][0]):
+        metadata = results["metadatas"][0][i]
+        filename = metadata["filename"]
+        formatted_results.append(f"Datei: {filename}\nText: {doc}\n")
+    
+    return "\n".join(formatted_results)
+
+def ask_llm(prompt):
+    result = client_llm.predict(
+        param_0=prompt,
+        api_name="/predict"
+    )
+    return result
+
 # Erstelle die Gradio-Schnittstelle
 with gr.Blocks() as demo:
-    gr.Markdown("# PDF Upload to ChromaDB")
-    file_input = gr.File(label="Wähle eine PDF-Datei aus", type="filepath")
-    submit_button = gr.Button("Hochladen")
-    output_text = gr.Textbox(label="Status")
-
-    submit_button.click(process_pdf, inputs=file_input, outputs=output_text)
+    gr.Markdown("# PDF Upload and Similarity Search with ChromaDB and LLM")
+    
+    with gr.Row():
+        file_input = gr.File(label="Wähle eine PDF-Datei aus", type="filepath")
+        submit_button = gr.Button("Hochladen")
+        upload_output = gr.Textbox(label="Upload Status")
+    
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Suche nach ähnlichen Dokumenten", placeholder="Gib einen Suchbegriff ein")
+        search_button = gr.Button("Suchen")
+        search_output = gr.Textbox(label="Ähnliche Dokumente")
+    
+    with gr.Row():
+        llm_prompt_input = gr.Textbox(label="Frage an das LLM", placeholder="Gib eine Frage ein")
+        llm_submit_button = gr.Button("Fragen")
+        llm_output = gr.Textbox(label="LLM Antwort")
+    
+    submit_button.click(process_pdf, inputs=file_input, outputs=upload_output)
+    search_button.click(search_similar_documents, inputs=prompt_input, outputs=search_output)
+    llm_submit_button.click(ask_llm, inputs=llm_prompt_input, outputs=llm_output)
 
 # Starte die Gradio-Anwendung
 demo.launch()