Spaces:
Runtime error
Runtime error
File size: 4,960 Bytes
5922680 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | import os
import uuid
import json
import pickle
import concurrent.futures
from typing import List, Dict
from unstructured.partition.pdf import partition_pdf
from unstructured.chunking.title import chunk_by_title
from langchain_core.documents import Document
from langchain_pinecone import PineconeVectorStore
from langchain_core.messages import HumanMessage
from pinecone import Pinecone
from config import *
# Initialize Pinecone Client
if PINECONE_API_KEY:
pc = Pinecone(api_key=PINECONE_API_KEY)
else:
pc = None
class SessionDocStore:
"""
In-memory storage for heavy content (Images/Tables) tied to a session/document.
"""
def __init__(self):
self.store = {}
def save_chunk(self, doc_id: str, data: Dict):
self.store[doc_id] = data
def get_chunk(self, doc_id: str):
return self.store.get(doc_id, {})
def clear(self):
self.store = {}
def save_to_disk(self, file_path: str):
with open(file_path, 'wb') as f:
pickle.dump(self.store, f)
def load_from_disk(self, file_path: str):
if os.path.exists(file_path):
with open(file_path, 'rb') as f:
self.store = pickle.load(f)
def create_multimodal_summary(text, tables, images, api_key: str):
llm = get_llm(api_key=api_key)
prompt_text = f"Analyze content. TEXT: {text[:1000]}. INSTRUCTIONS: Summarize text and describe tables in detail without missing out /skipping any information."
message_content = [{"type": "text", "text": prompt_text}]
if images:
for b64_str in images:
if "," in b64_str: b64_str = b64_str.split(",")[1]
message_content.append({
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{b64_str}", "detail": "low"}
})
response = llm.invoke([HumanMessage(content=message_content)])
return response.content
def process_single_chunk(i, chunk, doc_store, api_key: str):
"""
Worker function to process a single chunk in a separate thread.
"""
content = {'text': chunk.text, 'tables': [], 'images': []}
# Extract visual data
if hasattr(chunk, 'metadata') and hasattr(chunk.metadata, 'orig_elements'):
for element in chunk.metadata.orig_elements:
el_type = type(element).__name__
if el_type == 'Table':
content['tables'].append(getattr(element.metadata, 'text_as_html', element.text))
elif el_type == 'Image' and hasattr(element.metadata, 'image_base64'):
content['images'].append(element.metadata.image_base64)
# Deciding whether to call LLM (Slow) or just use text (Fast)
if content['images'] or content['tables']:
enhanced_text = create_multimodal_summary(content['text'], content['tables'], content['images'], api_key)
else:
enhanced_text = content['text']
doc_id = str(uuid.uuid4())
# Save heavy data to local store
doc_store.save_chunk(doc_id, {
"raw_text": content['text'],
"tables": content['tables'],
"images": content['images']
})
# Return the processed Document
return Document(
page_content=enhanced_text,
metadata={"doc_id": doc_id, "chunk_index": i}
)
def process_and_ingest(file_path: str, doc_store: SessionDocStore, api_key: str):
if not pc:
raise ValueError("PINECONE_API_KEY is not set.")
print(f"📄 Partitioning: {file_path}")
elements = partition_pdf(
filename=file_path,
strategy="auto",
infer_table_structure=True,
extract_image_block_types=["Table"],
extract_image_block_to_payload=True
)
chunks = chunk_by_title(elements, max_characters=2000, new_after_n_chars=1500, combine_text_under_n_chars=300)
documents = []
print(f"🔄 Processing {len(chunks)} chunks in parallel...")
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
future_to_chunk = {
executor.submit(process_single_chunk, i, chunk, doc_store, api_key): i
for i, chunk in enumerate(chunks)
}
futures_list = list(future_to_chunk.keys())
for future in futures_list:
try:
doc = future.result()
documents.append(doc)
except Exception as e:
print(f"❌ Error processing chunk: {e}")
print(f"🔮 Ingesting {len(documents)} vectors to Pinecone Index: {INDEX_NAME}")
# Clear index first if necessary, or just use a specific namespace. For this single-file app, we can use a hardcoded namespace
namespace = "single-pdf-mcp"
PineconeVectorStore.from_documents(
documents=documents,
index_name=INDEX_NAME,
embedding=get_embeddings(),
namespace=namespace
)
return documents
|