Update helper/process_pdf.py

helper/process_pdf.py

@@ -1,347 +1,351 @@
-import requests
-from pdf2image import convert_from_path
-import base64
-from pymongo import MongoClient
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_google_genai import GoogleGenerativeAIEmbeddings
-from langchain_core.messages import HumanMessage
-import os
-import re
-import json
-import uuid
-from dotenv import load_dotenv
-import pinecone
-load_dotenv()
-MONGO_URI = os.getenv("MONGO_URI")
-DB_NAME = os.getenv("DB_NAME")
-COLLECTION_NAME = os.getenv("COLLECTION_NAME")
-FLASH_API = os.getenv("FLASH_API")
-mongo_client = MongoClient(MONGO_URI)
-db = mongo_client[DB_NAME]
-collection = db[COLLECTION_NAME]
-collection2=db['about_company']
-model = ChatGoogleGenerativeAI(model="gemini-1.5-flash-002", temperature=0.2, max_tokens=None, google_api_key=FLASH_API)
-google_embeddings = GoogleGenerativeAIEmbeddings(
-    model="models/embedding-001",  # Correct model name
-    google_api_key="AIzaSyANNRKfEb-YnVIBaSAq6hQ38XpxxGwvaws"  # Your API key
-)
-pc = pinecone.Pinecone(
-    api_key="4a80f293-ae6d-489a-a7d8-33ea3fcdd26b"  # Your Pinecone API key
-)
-index_name = "mospi"
-index = pc.Index(index_name)
-
-about_company_doc=collection2.find_one({"type":"about_company"})
-if about_company_doc:
-    about_company=about_company_doc.get('company_description','')
-
-pdf_temp_dir = 'temp/pdf_files'
-image_temp_dir = 'temp/page_images'
-
-os.makedirs(pdf_temp_dir, exist_ok=True)
-os.makedirs(image_temp_dir, exist_ok=True)
-
-pdf_path = os.path.join(pdf_temp_dir, 'downloaded_file.pdf')
-
-def download_and_split_pdf_to_image(url):
-    try:
-        response = requests.get(url)
-        with open(pdf_path, 'wb') as pdf_file:
-            pdf_file.write(response.content)
-
-
-    except Exception as e:
-        print(f"error occured during downloading pdf from object url : {e}")
-        return None
-
-    try:
-        images = convert_from_path(pdf_path)
-        for i, image in enumerate(images):
-            image_path = os.path.join(image_temp_dir, f'page_{i + 1}.png')
-            image.save(image_path, 'PNG')
-            print(f'Saved image: {image_path}')
-        return True
-
-    except Exception as e:
-        print(f"error occured in converting pdf pages to image : {e}")
-        return None
-
-
-
-system_prompt_text = f"""Given is an image of a PDF page.Your task is to extract all the information from this image and give a detailed summary of the page, do not miss out on any information, include keywords or any terms mentioned in the pdf.Its possible that sometimes the information might look redundant or repetitive still extract it.'
-Given below is a company information whose pdf page is givn to you, to understand the context.
-- About Company: {about_company}
-Follow this Expected output format given below:
-Expected Output format : {{"description":"String"}}
-
-"""
-
-
-
-# def process_image_using_llm(image, page_number, url):
-#     try:
-#         message = HumanMessage(
-#             content=[
-#                 {"type": "text", "text": system_prompt_text},
-#                 {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}},
-#             ],
-#         )
-#         response = model.invoke([message])
-#         print(f"LLM response for page {page_number}: {response}")
-#
-#         # Extract JSON from the response content using regex
-#         match = re.search(r"\{.*\}", response.content.strip())
-#         if match:
-#             json_data = match.group(0)
-#
-#             # Step 1: Locate the "description" field and escape all single quotes within it
-#             description_match = re.search(r"'description'\s*:\s*('.*?'|\".*?\")", json_data)
-#
-#             if description_match:
-#                 description_text = description_match.group(1)
-#
-#                 # Replace outer single quotes with double quotes if necessary
-#                 if description_text.startswith("'") and description_text.endswith("'"):
-#                     description_text = f'"{description_text[1:-1]}"'
-#                 elif description_text.startswith('"') and description_text.endswith('"'):
-#                     pass  # No change needed if already double quotes
-#
-#                 # Escape all single quotes within the description text
-#                 description_text = description_text.replace("'", "\\'")
-#
-#                 # Replace the original match with the updated description text
-#                 json_data = (
-#                     json_data[:description_match.start(1)] +
-#                     description_text +
-#                     json_data[description_match.end(1):]
-#                 )
-#
-#             # Step 2: Attempt to load the cleaned JSON string
-#             try:
-#                 data = json.loads(json_data)  # Load as JSON
-#                 description = data.get("description", "None").strip()
-#                 can_find_description = description != "None"
-#
-#                 return {
-#                     "page_number": page_number,
-#                     "description": description if can_find_description else None,
-#                     "can_find_description": can_find_description
-#                 }
-#             except json.JSONDecodeError as e:
-#                 print(f"Error decoding JSON for page {page_number}: {e}")
-#                 return {
-#                     "page_number": page_number,
-#                     "description": None,
-#                     "can_find_description": False
-#                 }
-#         else:
-#             print(f"No valid JSON found in the response for page {page_number}")
-#             return {
-#                 "page_number": page_number,
-#                 "description": None,
-#                 "can_find_description": False
-#             }
-#
-#     except Exception as e:
-#         print(f"Error processing page {page_number}: {e}")
-#         return {
-#             "page_number": page_number,
-#             "description": None,
-#             "can_find_description": False
-#         }
-#
-
-def process_image_using_llm(image, page_number, url, max_retries=3):
-    for attempt in range(1, max_retries + 1):
-        try:
-            # Send the image and system prompt to the LLM
-            message = HumanMessage(
-                content=[
-                    {"type": "text", "text": system_prompt_text},
-                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}},
-                ],
-            )
-            response = model.invoke([message])
-            print(f"Attempt {attempt}: LLM response for page {page_number}: {response}")
-
-            # Retry only if JSON data is missing or JSON decode error occurs
-            try:
-                # Check for JSON content in the response
-                match = re.search(r"\{.*\}", response.content.strip())
-                if match:
-                    json_data = match.group(0)
-
-                    # Locate the "description" field and escape single quotes within it
-                    description_match = re.search(r"'description'\s*:\s*('.*?'|\".*?\")", json_data)
-                    if description_match:
-                        description_text = description_match.group(1)
-                        if description_text.startswith("'") and description_text.endswith("'"):
-                            description_text = f'"{description_text[1:-1]}"'
-                        description_text = description_text.replace("'", "\\'")
-
-                        # Replace the original match with the updated description text
-                        json_data = (
-                                json_data[:description_match.start(1)] +
-                                description_text +
-                                json_data[description_match.end(1):]
-                        )
-
-                    # Load the cleaned JSON data
-                    data = json.loads(json_data)
-                    description = data.get("description", "None").strip()
-                    can_find_description = description != "None"
-
-                    return {
-                        "page_number": page_number,
-                        "description": description if can_find_description else None,
-                        "can_find_description": can_find_description
-                    }
-                else:
-                    print(f"No valid JSON found in the response for page {page_number} on attempt {attempt}")
-
-                # Trigger retry if no JSON data is found
-                if attempt == max_retries:
-                    print(f"Max retries reached for page {page_number}. Skipping.")
-                    return {
-                        "page_number": page_number,
-                        "description": None,
-                        "can_find_description": False
-                    }
-
-            except json.JSONDecodeError as e:
-                print(f"JSON decode error on attempt {attempt} for page {page_number}: {e}")
-                if attempt == max_retries:
-                    return {
-                        "page_number": page_number,
-                        "description": None,
-                        "can_find_description": False
-                    }
-
-        # Handle any other exceptions without retrying
-        except Exception as e:
-            print(f"Outer exception for page {page_number}: {e}")
-            return {
-                "page_number": page_number,
-                "description": None,
-                "can_find_description": False
-            }
-
-
-def create_embedding_for_pdf_chunks(page,description,url,tags,categories):
-    try:
-        document = collection.find_one({'object_url': url})
-        file_type = document.get("type")
-        mongo_id = str(document.get('_id'))
-        embedding = google_embeddings.embed_query(description)
-        pinecone_id = str(uuid.uuid4())
-
-        vectors = [{
-            'id': pinecone_id,
-            'values': embedding,
-            'metadata': {
-                'description': description,
-                "url": url,
-                "page_number":page,
-                "tag": file_type,
-                "mongo_id": mongo_id,
-                "tags": ','.join(tags),
-                "categories": ','.join(categories)  # Store MongoDB ID in metadata
-            }
-        }]
-        index.upsert(vectors)
-        print(f"Inserted: page {page} in Pinecone with MongoDB ID {mongo_id} in metadata")
-
-        collection.update_one(
-            {
-                "_id": document["_id"],
-                "chunks.page_number": page  # Match document and specific chunk by page number
-            },
-            {
-                "$set": {
-                    "chunks.$.pinecone_id": pinecone_id,
-                    "chunks.$.successfully_embedding_created": True
-                }
-            }
-        )
-        return True
-
-    except Exception as e:
-        print(f"error occured in creating embedding for pdf with mongo id {mongo_id} for page {page}")
-        collection.update_one(
-            {
-                "_id": document["_id"],
-                "chunks.page_number": page  # Match document and specific chunk by page number
-            },
-            {
-                "$set": {
-                    "chunks.$.successfully_embedding_created": False
-                }
-            }
-        )
-        return False
-
-
-
-def cleanup_directory(directory_path):
-    try:
-        for filename in os.listdir(directory_path):
-            file_path = os.path.join(directory_path, filename)
-            if os.path.isfile(file_path):
-                os.remove(file_path)
-        print(f"Cleaned up files in {directory_path}")
-    except Exception as e:
-        print(f"Error cleaning up directory {directory_path}: {e}")
-
-
-def process_pdf(url,tags,categories):
-    print(f"Processing PDF with URL: {url}")
-    if download_and_split_pdf_to_image(url):
-        chunks = []
-        image_files = sorted(
-            os.listdir(image_temp_dir),
-            key=lambda x: int(re.search(r'page_(\d+)', x).group(1))
-        )
-        for count, image_name in enumerate(image_files, start=1):
-            print(f"Processing page {count} of the PDF")
-            image_path = os.path.join(image_temp_dir, image_name)
-            with open(image_path, "rb") as image_file:
-                image_data = base64.b64encode(image_file.read()).decode("utf-8")
-                page_result = process_image_using_llm(image_data,count, url)
-                chunks.append(page_result)
-
-        # Update the MongoDB document with the chunks array
-        collection.update_one(
-            {"object_url": url},
-            {"$set": {"chunks": chunks}},
-            upsert=True
-        )
-        print("Saved chunks to MongoDB.")
-
-        total_pages=len(chunks)
-        failed_pages=0
-        for chunk in chunks:
-            page_number=chunk.get('page_number')
-            description=chunk.get("description")
-            if description:
-                embedding_created=create_embedding_for_pdf_chunks(page_number,description,url,tags,categories)
-                if not embedding_created:
-                    failed_pages+=1
-            else:
-                print(f"skipping the page {page_number} as description is None")
-                failed_pages+=1
-                continue
-
-        cleanup_directory(pdf_temp_dir)
-        cleanup_directory(image_temp_dir)
-
-        return failed_pages < total_pages
-
-
-
-
-
-
-
-
-
-
+import requests
+from pdf2image import convert_from_path
+import base64
+from pymongo import MongoClient
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_core.messages import HumanMessage
+import os
+import re
+import json
+import uuid
+from dotenv import load_dotenv
+import pinecone
+load_dotenv()
+MONGO_URI = os.getenv("MONGO_URI")
+DB_NAME = os.getenv("DB_NAME")
+COLLECTION_NAME = os.getenv("COLLECTION_NAME")
+FLASH_API = os.getenv("FLASH_API")
+PINECONE_API=os.getenv("PINECONE_API")
+PINECONE_INDEX=os.getenv("PINECONE_INDEX")
+
+
+mongo_client = MongoClient(MONGO_URI)
+db = mongo_client[DB_NAME]
+collection = db[COLLECTION_NAME]
+collection2=db['about_company']
+model = ChatGoogleGenerativeAI(model="gemini-1.5-flash-002", temperature=0.2, max_tokens=None, google_api_key=FLASH_API)
+google_embeddings = GoogleGenerativeAIEmbeddings(
+    model="models/embedding-001",  # Correct model name
+    google_api_key=FLASH_API  # Your API key
+)
+pc = pinecone.Pinecone(
+    api_key=PINECONE_API  # Your Pinecone API key
+)
+
+index = pc.Index(PINECONE_INDEX)
+
+about_company_doc=collection2.find_one({"type":"about_company"})
+if about_company_doc:
+    about_company=about_company_doc.get('company_description','')
+
+pdf_temp_dir = 'temp/pdf_files'
+image_temp_dir = 'temp/page_images'
+
+os.makedirs(pdf_temp_dir, exist_ok=True)
+os.makedirs(image_temp_dir, exist_ok=True)
+
+pdf_path = os.path.join(pdf_temp_dir, 'downloaded_file.pdf')
+
+def download_and_split_pdf_to_image(url):
+    try:
+        response = requests.get(url)
+        with open(pdf_path, 'wb') as pdf_file:
+            pdf_file.write(response.content)
+
+
+    except Exception as e:
+        print(f"error occured during downloading pdf from object url : {e}")
+        return None
+
+    try:
+        images = convert_from_path(pdf_path)
+        for i, image in enumerate(images):
+            image_path = os.path.join(image_temp_dir, f'page_{i + 1}.png')
+            image.save(image_path, 'PNG')
+            print(f'Saved image: {image_path}')
+        return True
+
+    except Exception as e:
+        print(f"error occured in converting pdf pages to image : {e}")
+        return None
+
+
+
+system_prompt_text = f"""Given is an image of a PDF page.Your task is to extract all the information from this image and give a detailed summary of the page, do not miss out on any information, include keywords or any terms mentioned in the pdf.Its possible that sometimes the information might look redundant or repetitive still extract it.'
+Given below is a company information whose pdf page is givn to you, to understand the context.
+- About Company: {about_company}
+Follow this Expected output format given below:
+Expected Output format : {{"description":"String"}}
+
+"""
+
+
+
+# def process_image_using_llm(image, page_number, url):
+#     try:
+#         message = HumanMessage(
+#             content=[
+#                 {"type": "text", "text": system_prompt_text},
+#                 {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}},
+#             ],
+#         )
+#         response = model.invoke([message])
+#         print(f"LLM response for page {page_number}: {response}")
+#
+#         # Extract JSON from the response content using regex
+#         match = re.search(r"\{.*\}", response.content.strip())
+#         if match:
+#             json_data = match.group(0)
+#
+#             # Step 1: Locate the "description" field and escape all single quotes within it
+#             description_match = re.search(r"'description'\s*:\s*('.*?'|\".*?\")", json_data)
+#
+#             if description_match:
+#                 description_text = description_match.group(1)
+#
+#                 # Replace outer single quotes with double quotes if necessary
+#                 if description_text.startswith("'") and description_text.endswith("'"):
+#                     description_text = f'"{description_text[1:-1]}"'
+#                 elif description_text.startswith('"') and description_text.endswith('"'):
+#                     pass  # No change needed if already double quotes
+#
+#                 # Escape all single quotes within the description text
+#                 description_text = description_text.replace("'", "\\'")
+#
+#                 # Replace the original match with the updated description text
+#                 json_data = (
+#                     json_data[:description_match.start(1)] +
+#                     description_text +
+#                     json_data[description_match.end(1):]
+#                 )
+#
+#             # Step 2: Attempt to load the cleaned JSON string
+#             try:
+#                 data = json.loads(json_data)  # Load as JSON
+#                 description = data.get("description", "None").strip()
+#                 can_find_description = description != "None"
+#
+#                 return {
+#                     "page_number": page_number,
+#                     "description": description if can_find_description else None,
+#                     "can_find_description": can_find_description
+#                 }
+#             except json.JSONDecodeError as e:
+#                 print(f"Error decoding JSON for page {page_number}: {e}")
+#                 return {
+#                     "page_number": page_number,
+#                     "description": None,
+#                     "can_find_description": False
+#                 }
+#         else:
+#             print(f"No valid JSON found in the response for page {page_number}")
+#             return {
+#                 "page_number": page_number,
+#                 "description": None,
+#                 "can_find_description": False
+#             }
+#
+#     except Exception as e:
+#         print(f"Error processing page {page_number}: {e}")
+#         return {
+#             "page_number": page_number,
+#             "description": None,
+#             "can_find_description": False
+#         }
+#
+
+def process_image_using_llm(image, page_number, url, max_retries=3):
+    for attempt in range(1, max_retries + 1):
+        try:
+            # Send the image and system prompt to the LLM
+            message = HumanMessage(
+                content=[
+                    {"type": "text", "text": system_prompt_text},
+                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}},
+                ],
+            )
+            response = model.invoke([message])
+            print(f"Attempt {attempt}: LLM response for page {page_number}: {response}")
+
+            # Retry only if JSON data is missing or JSON decode error occurs
+            try:
+                # Check for JSON content in the response
+                match = re.search(r"\{.*\}", response.content.strip())
+                if match:
+                    json_data = match.group(0)
+
+                    # Locate the "description" field and escape single quotes within it
+                    description_match = re.search(r"'description'\s*:\s*('.*?'|\".*?\")", json_data)
+                    if description_match:
+                        description_text = description_match.group(1)
+                        if description_text.startswith("'") and description_text.endswith("'"):
+                            description_text = f'"{description_text[1:-1]}"'
+                        description_text = description_text.replace("'", "\\'")
+
+                        # Replace the original match with the updated description text
+                        json_data = (
+                                json_data[:description_match.start(1)] +
+                                description_text +
+                                json_data[description_match.end(1):]
+                        )
+
+                    # Load the cleaned JSON data
+                    data = json.loads(json_data)
+                    description = data.get("description", "None").strip()
+                    can_find_description = description != "None"
+
+                    return {
+                        "page_number": page_number,
+                        "description": description if can_find_description else None,
+                        "can_find_description": can_find_description
+                    }
+                else:
+                    print(f"No valid JSON found in the response for page {page_number} on attempt {attempt}")
+
+                # Trigger retry if no JSON data is found
+                if attempt == max_retries:
+                    print(f"Max retries reached for page {page_number}. Skipping.")
+                    return {
+                        "page_number": page_number,
+                        "description": None,
+                        "can_find_description": False
+                    }
+
+            except json.JSONDecodeError as e:
+                print(f"JSON decode error on attempt {attempt} for page {page_number}: {e}")
+                if attempt == max_retries:
+                    return {
+                        "page_number": page_number,
+                        "description": None,
+                        "can_find_description": False
+                    }
+
+        # Handle any other exceptions without retrying
+        except Exception as e:
+            print(f"Outer exception for page {page_number}: {e}")
+            return {
+                "page_number": page_number,
+                "description": None,
+                "can_find_description": False
+            }
+
+
+def create_embedding_for_pdf_chunks(page,description,url,tags,categories):
+    try:
+        document = collection.find_one({'object_url': url})
+        file_type = document.get("type")
+        mongo_id = str(document.get('_id'))
+        embedding = google_embeddings.embed_query(description)
+        pinecone_id = str(uuid.uuid4())
+
+        vectors = [{
+            'id': pinecone_id,
+            'values': embedding,
+            'metadata': {
+                'description': description,
+                "url": url,
+                "page_number":page,
+                "tag": file_type,
+                "mongo_id": mongo_id,
+                "tags": ','.join(tags),
+                "categories": ','.join(categories)  # Store MongoDB ID in metadata
+            }
+        }]
+        index.upsert(vectors)
+        print(f"Inserted: page {page} in Pinecone with MongoDB ID {mongo_id} in metadata")
+
+        collection.update_one(
+            {
+                "_id": document["_id"],
+                "chunks.page_number": page  # Match document and specific chunk by page number
+            },
+            {
+                "$set": {
+                    "chunks.$.pinecone_id": pinecone_id,
+                    "chunks.$.successfully_embedding_created": True
+                }
+            }
+        )
+        return True
+
+    except Exception as e:
+        print(f"error occured in creating embedding for pdf with mongo id {mongo_id} for page {page}")
+        collection.update_one(
+            {
+                "_id": document["_id"],
+                "chunks.page_number": page  # Match document and specific chunk by page number
+            },
+            {
+                "$set": {
+                    "chunks.$.successfully_embedding_created": False
+                }
+            }
+        )
+        return False
+
+
+
+def cleanup_directory(directory_path):
+    try:
+        for filename in os.listdir(directory_path):
+            file_path = os.path.join(directory_path, filename)
+            if os.path.isfile(file_path):
+                os.remove(file_path)
+        print(f"Cleaned up files in {directory_path}")
+    except Exception as e:
+        print(f"Error cleaning up directory {directory_path}: {e}")
+
+
+def process_pdf(url,tags,categories):
+    print(f"Processing PDF with URL: {url}")
+    if download_and_split_pdf_to_image(url):
+        chunks = []
+        image_files = sorted(
+            os.listdir(image_temp_dir),
+            key=lambda x: int(re.search(r'page_(\d+)', x).group(1))
+        )
+        for count, image_name in enumerate(image_files, start=1):
+            print(f"Processing page {count} of the PDF")
+            image_path = os.path.join(image_temp_dir, image_name)
+            with open(image_path, "rb") as image_file:
+                image_data = base64.b64encode(image_file.read()).decode("utf-8")
+                page_result = process_image_using_llm(image_data,count, url)
+                chunks.append(page_result)
+
+        # Update the MongoDB document with the chunks array
+        collection.update_one(
+            {"object_url": url},
+            {"$set": {"chunks": chunks}},
+            upsert=True
+        )
+        print("Saved chunks to MongoDB.")
+
+        total_pages=len(chunks)
+        failed_pages=0
+        for chunk in chunks:
+            page_number=chunk.get('page_number')
+            description=chunk.get("description")
+            if description:
+                embedding_created=create_embedding_for_pdf_chunks(page_number,description,url,tags,categories)
+                if not embedding_created:
+                    failed_pages+=1
+            else:
+                print(f"skipping the page {page_number} as description is None")
+                failed_pages+=1
+                continue
+
+        cleanup_directory(pdf_temp_dir)
+        cleanup_directory(image_temp_dir)
+
+        return failed_pages < total_pages
+
+
+
+
+
+
+
+
+
+