Update src/utils/ingest_text.py

src/utils/ingest_text.py

@@ -1,120 +1,96 @@
+import os
+import pickle
+from typing import List
 from llama_parse import LlamaParse
-from langchain_chroma import Chroma
-from qdrant_client import QdrantClient
-from langchain_community.vectorstores.qdrant import Qdrant
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
 from langchain_community.document_loaders.directory import DirectoryLoader
-import os
-from fastembed import TextEmbedding
-from typing import List
+from langchain_community.document_loaders import TextLoader
+from langchain_community.vectorstores.qdrant import Qdrant
+from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
 import nltk
-nltk.download('punkt')
-
-
 import nest_asyncio
+
+# Setup
+nltk.download('punkt')
 nest_asyncio.apply()
 
+# Load environment variables
+from dotenv import load_dotenv
+load_dotenv()
+
+# Environment keys
 llamaparse_api_key = os.getenv("LLAMA_CLOUD_API_KEY")
-#qdrant_url = os.getenv("QDRANT_URL ")
-#qdrant_api_key = os.getenv("QDRANT_API_KEY")
 groq_api_key = os.getenv("GROQ_API_KEY")
 
+# Paths
+parsed_data_file = os.path.join("data", "parsed_data.pkl")
+output_md = os.path.join("data", "output.md")
+md_directory = "data"
+collection_name = "rag"
 
-parsed_data_file = r"./data/parsed_data.pkl"
-output_md = r"./data/output.md"
-loki = r"data"
-
-import pickle
-# Define a function to load parsed data if available, or parse if not
-def load_or_parse_data(loc):
-    data_file = parsed_data_file
-
-    if os.path.exists(data_file):
-        # Load the parsed data from the file
-        with open(data_file, "rb") as f:
+# Helper: Load or parse PDF
+def load_or_parse_data(pdf_path):
+    if os.path.exists(parsed_data_file):
+        with open(parsed_data_file, "rb") as f:
             parsed_data = pickle.load(f)
     else:
-        # Perform the parsing step and store the result in llama_parse_documents
-        parsingInstructiontest10k = """The provided document is an user guide or a manual.
-        It contains many images and tables.
-        Try to be precise while answering the questions"""
-        parser = LlamaParse(api_key=llamaparse_api_key, result_type="markdown", parsing_instruction=parsingInstructiontest10k) # type: ignore
-        llama_parse_documents = parser.load_data(loc)
-
+        parsing_instruction = """The provided document is a user guide or manual.
+        It contains many images and tables. Be precise while answering questions."""
+        parser = LlamaParse(api_key=llamaparse_api_key, result_type="markdown", parsing_instruction=parsing_instruction)  # type: ignore
+        parsed_data = parser.load_data(pdf_path)
 
-        # Save the parsed data to a file
-        with open(data_file, "wb") as f:
-            pickle.dump(llama_parse_documents, f)
-
-        # Set the parsed data to the variable
-        parsed_data = llama_parse_documents
+        with open(parsed_data_file, "wb") as f:
+            pickle.dump(parsed_data, f)
 
     return parsed_data
 
+# Main vector DB builder
+def create_vector_database(pdf_path):
+    print("🧠 Starting vector DB creation...")
 
-# Create vector database
-def create_vector_database(loc):
-    """
-    Creates a vector database using document loaders and embeddings.
+    parsed_docs = load_or_parse_data(pdf_path)
+    if not parsed_docs:
+        raise ValueError("❌ No parsed documents returned from LlamaParse!")
 
-    This function loads urls,
-    splits the loaded documents into chunks, transforms them into embeddings using OllamaEmbeddings,
-    and finally persists the embeddings into a Chroma vector database.
+    os.makedirs(md_directory, exist_ok=True)
 
-    """
-    # Call the function to either load or parse the data
-    print("text_db")
-    llama_parse_documents = load_or_parse_data(loc)
-    #print(llama_parse_documents[1].text[:100])
+    # Write Markdown content to file (overwrite)
+    with open(output_md, 'w', encoding='utf-8') as f:
+        for doc in parsed_docs:
+            if hasattr(doc, "text") and doc.text.strip():
+                f.write(doc.text.strip() + "\n\n")
 
-    #with open('data/output.md', 'a') as f:  # Open the file in append mode ('a')
-      #  for doc in llama_parse_documents:
-       #     f.write(doc.text + '\n')
-    with open(output_md,'a', encoding='utf-8') as f:  # Open the file in append mode ('a')
-        for doc in llama_parse_documents:
-            f.write(doc.text + '\n')
+    # Ensure .md file was written
+    if not os.path.exists(output_md) or os.path.getsize(output_md) == 0:
+        raise RuntimeError("❌ Markdown file was not created or is empty!")
 
-    loader = DirectoryLoader(loki, glob="**/*.md", show_progress=True)
-    documents = loader.load()
-    # Split loaded documents into chunks
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=100)
-    print('data chunckex')
-    docs = text_splitter.split_documents(documents)
-    print(len(docs))
+    # Load documents
+    try:
+        loader = DirectoryLoader(md_directory, glob="**/*.md", show_progress=True)
+        documents = loader.load()
+    except Exception as e:
+        print("⚠️ DirectoryLoader failed, falling back to TextLoader...")
+        documents = TextLoader(output_md, encoding='utf-8').load()
 
-    #len(docs)
-    #docs[0]
+    if not documents:
+        raise RuntimeError("❌ No documents loaded from markdown!")
 
-    # Initialize Embeddings
-    embeddings = FastEmbedEmbeddings() # type: ignore
-    #embeddings = TextEmbedding()
+    # Split documents
+    splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=100)
+    docs = splitter.split_documents(documents)
+    print(f"✅ Loaded and split {len(docs)} chunks.")
 
-    print('Vector DB started!')
+    # Embedding
+    embeddings = FastEmbedEmbeddings()  # type: ignore
 
-    # Create and persist a Chroma vector database from the chunked documents
+    # Create vector store
+    print("📦 Creating Qdrant vector DB...")
     qdrant = Qdrant.from_documents(
         documents=docs,
         embedding=embeddings,
-        path=r".\data\local_qdrant",
-        #url=qdrant_url,
-        collection_name="rag"
-        #api_key=qdrant_api_key
+        path=os.path.join("data", "local_qdrant"),
+        collection_name=collection_name,
     )
-    # save to disk
-    #db2 = Chroma.from_documents(docs, embeddings, persist_directory="./chroma_db")
-    #docs = db2.similarity_search(query)
-
-# load from disk
-    #db3 = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
-
-    #query it
-    #query = "what is the agend of Financial Statements for 2022 ?"
-    #found_doc = qdrant.similarity_search(query, k=3)
-    #print(found_doc[0][:100])
-    #
-    print('Vector DB created successfully !')
-    #query = "Switching between external devices connected to the TV"
-    #found_doc = qdrant.similarity_search(query, k=3)
-    #print(found_doc)
-    return qdrant
+
+    print("✅ Vector DB created successfully.")
+    return qdrant