Update modules/analysis.py

modules/analysis.py

@@ -1,52 +1,37 @@
-from langchain.chains import RetrievalQA
-from langchain.llms import HuggingFacePipeline
-from transformers import pipeline
-from modules import parser, vectorizer
-
-def run_analysis(uploaded_files, text_input, query, quick_action, temperature, start_time, end_time):
-    """
-    Main logic that runs when the user clicks 'Analyze Logs'.
-    It combines file and text inputs, applies embeddings,
-    and performs question answering using a language model.
-    """
-    logs_text = ""
-
-    # Combine all uploaded files into one text string
-    if uploaded_files:
-        logs_text += parser.parse_uploaded_files(uploaded_files)
-
-    # Add manual pasted text logs
-    if text_input:
-        logs_text += "\n" + text_input
-
-    # Show error if no log input provided
-    if not logs_text.strip():
-        return "❌ No logs provided.", None, None, None
-
-    # Use either free-form query or a quick action
-    query_text = query if query else quick_action
-    if not query_text:
-        return "❌ No query provided.", None, None, None
-
-    # Chunk logs and embed them
-    docs = vectorizer.prepare_documents(logs_text)
-    vectordb = vectorizer.create_vectorstore(docs)
-
-    # Load a small Hugging Face text generation pipeline (GPT-2 here)
-    pipe = pipeline("text-generation", model="gpt2", max_length=512, temperature=temperature)
-    llm = HuggingFacePipeline(pipeline=pipe)
-
-    # Create LangChain retrieval-based QA chain
-    qa = RetrievalQA.from_chain_type(llm=llm, retriever=vectordb.as_retriever())
-
-    # Run the query against embedded document chunks
-    result = qa.run(query_text)
-
-    # -------- Mocked example chart and alert outputs --------
-    bar_data = {"Hour": ["14:00", "15:00"], "Count": [8, 4]}
-    pie_data = {"Event Type": ["Blocked", "Scan"], "Count": [8, 4]}
-    alerts = [("CRITICAL", "8 blocked SSH attempts from 192.168.1.5"),
-              ("WARNING", "4 port scanning alerts from 10.0.0.8")]
-
-    # Return structured outputs to Gradio UI
-    return result, bar_data, pie_data, alerts
+from langchain.chains import RetrievalQA
+from langchain_community.llms import HuggingFacePipeline  # ✅ Updated import
+from transformers import pipeline
+from modules import parser, vectorizer
+
+def run_analysis(uploaded_files, text_input, query, quick_action, temperature, start_time, end_time):
+    logs_text = ""
+
+    if uploaded_files:
+        logs_text += parser.parse_uploaded_files(uploaded_files)
+
+    if text_input:
+        logs_text += "\n" + text_input
+
+    if not logs_text.strip():
+        return "❌ No logs provided.", None, None, None
+
+    query_text = query if query else quick_action
+    if not query_text:
+        return "❌ No query provided.", None, None, None
+
+    docs = vectorizer.prepare_documents(logs_text)
+    vectordb = vectorizer.create_vectorstore(docs)
+
+    pipe = pipeline("text-generation", model="gpt2", max_length=512, temperature=temperature)
+    llm = HuggingFacePipeline(pipeline=pipe)
+
+    qa = RetrievalQA.from_chain_type(llm=llm, retriever=vectordb.as_retriever())
+    result = qa.run(query_text)
+
+    # Example mock output
+    bar_data = {"Hour": ["14:00", "15:00"], "Count": [8, 4]}
+    pie_data = {"Event Type": ["Blocked", "Scan"], "Count": [8, 4]}
+    alerts = [("CRITICAL", "8 blocked SSH attempts from 192.168.1.5"),
+              ("WARNING", "4 port scanning alerts from 10.0.0.8")]
+
+    return result, bar_data, pie_data, alerts