Update rag_agent.py

rag_agent.py

@@ -1,71 +1,93 @@
-import os
-from langchain_community.document_loaders import PyPDFLoader, TextLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import FAISS
-from langchain.prompts import PromptTemplate
-from langchain_core.output_parsers import StrOutputParser
-from llm_setup import llm, embeddings
-
-if not llm or not embeddings:
-    raise ImportError("LLM or Embedding models could not be loaded. Please check llm_setup.py.")
-
-def get_document_context(file_path: str, query: str) -> str:
-    """
-    Loads a document, splits it, creates an in-memory FAISS vector store,
-    and retrieves the most relevant context for a given query.
-    """
-    print("--- Using FAISS for document retrieval ---")
-    
-    if file_path.endswith(".pdf"):
-        loader = PyPDFLoader(file_path)
-    elif file_path.endswith(".txt"):
-        loader = TextLoader(file_path)
-    else:
-        return "Error: Unsupported file format. Please upload a .pdf or .txt file."
-
-    try:
-        documents = loader.load()
-        if not documents:
-            return "Error: Document is empty or could not be read."
-            
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
-        docs = text_splitter.split_documents(documents)
-        
-        if not docs:
-            return "Error: Could not extract text chunks from the document."
-
-        db = FAISS.from_documents(docs, embeddings)
-        retriever = db.as_retriever(search_kwargs={'k': 3})
-        retrieved_docs = retriever.invoke(query)
-        
-        context = "\n\n".join([doc.page_content for doc in retrieved_docs])
-        return context
-    except Exception as e:
-        print(f"An error occurred during document processing: {e}")
-        return "Error: Failed to process the provided document."
-
-
-def run_rag_agent(user_prompt: str, file_path: str) -> str:
-    """
-    The main agentic function to retrieve context from a user-provided document.
-    """
-    print("--- RAG Agent Activated (Document-Only Mode) ---")
-    
-    prompt_template = PromptTemplate.from_template(
-        """You are a research assistant. Based on the user's story idea, what is the single most 
-        important keyword or question to search for within their provided document to find relevant context?
-        
-        User's Story Idea: '{prompt}'
-        
-        Optimized Search Query for Document:"""
-    )
-    
-    chain = prompt_template | llm | StrOutputParser()
-    
-    search_query = chain.invoke({"prompt": user_prompt})
-    print(f"Generated Search Query: {search_query}")
-
-    context = get_document_context(file_path, search_query)
-        
-    print("--- RAG Agent Finished ---")
-    return context
+import os
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+import PyPDF2
+from langchain_google_genai import ChatGoogleGenerativeAI
+import streamlit as st
+
+def load_document(file_path: str) -> str:
+    """Load document content from PDF or TXT file."""
+    try:
+        if file_path.endswith(".pdf"):
+            with open(file_path, 'rb') as file:
+                pdf_reader = PyPDF2.PdfReader(file)
+                text = ""
+                for page in pdf_reader.pages:
+                    text += page.extract_text() + "\n"
+                return text
+        elif file_path.endswith(".txt"):
+            with open(file_path, 'r', encoding='utf-8') as file:
+                return file.read()
+        else:
+            return "Error: Unsupported file format."
+    except Exception as e:
+        return f"Error reading file: {str(e)}"
+
+def simple_text_search(query: str, document_text: str, max_chunks: int = 3) -> str:
+    """Simple TF-IDF based text retrieval - much faster than FAISS."""
+    
+    # Split document into chunks
+    chunks = []
+    words = document_text.split()
+    chunk_size = 200  # words per chunk
+    
+    for i in range(0, len(words), chunk_size):
+        chunk = " ".join(words[i:i + chunk_size])
+        if chunk.strip():
+            chunks.append(chunk)
+    
+    if not chunks:
+        return "No content found in document."
+    
+    # Create TF-IDF vectors
+    vectorizer = TfidfVectorizer(stop_words='english', max_features=1000)
+    
+    try:
+        # Vectorize chunks and query
+        chunk_vectors = vectorizer.fit_transform(chunks)
+        query_vector = vectorizer.transform([query])
+        
+        # Calculate similarity
+        similarities = cosine_similarity(query_vector, chunk_vectors).flatten()
+        
+        # Get top matching chunks
+        top_indices = similarities.argsort()[-max_chunks:][::-1]
+        
+        relevant_chunks = [chunks[i] for i in top_indices if similarities[i] > 0.1]
+        
+        return "\n\n".join(relevant_chunks[:max_chunks])
+        
+    except Exception as e:
+        return f"Search error: {str(e)}"
+
+def run_rag_agent(user_prompt: str, file_path: str) -> str:
+    """Simple but effective RAG implementation."""
+    print("--- RAG Agent Activated (Lightweight Version) ---")
+    
+    # Load document
+    document_text = load_document(file_path)
+    if document_text.startswith("Error"):
+        return document_text
+    
+    # Generate search query using LLM
+    api_key = st.secrets.get("GEMINI_API_KEY", os.getenv("GEMINI_API_KEY"))
+    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", google_api_key=api_key)
+    
+    search_prompt = f"""Based on this story idea: "{user_prompt}"
+    
+What are the 2-3 most important keywords to search for in a document to find relevant context?
+Respond with just the keywords separated by spaces."""
+
+    try:
+        response = llm.invoke(search_prompt)
+        search_query = response.content.strip()
+        print(f"Generated Search Query: {search_query}")
+    except:
+        search_query = user_prompt  # Fallback
+    
+    # Retrieve relevant content
+    context = simple_text_search(search_query, document_text)
+    
+    print("--- RAG Agent Finished ---")
+    return context