File size: 3,392 Bytes
d9cd5c8 | 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 | import os
import requests
import tempfile
from typing import List
from langchain_community.document_loaders import PyPDFLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from langchain_core.tools import Tool
from pinecone import Pinecone, ServerlessSpec
from docx import Document as DocxDocument
from dotenv import load_dotenv
load_dotenv()
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
PINECONE_INDEX_NAME = os.getenv("document-processor-index-1") or "policy-index"
pc = Pinecone(api_key=PINECONE_API_KEY)
embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
def download_file(url: str) -> str:
try:
response = requests.get(url)
response.raise_for_status()
suffix = ".pdf" if ".pdf" in url else ".docx"
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
temp_file.write(response.content)
return temp_file.name
except requests.RequestException as e:
print(f"Error downloading file from {url}: {e}")
raise
def load_docx_documents(file_path: str) -> List[Document]:
doc = DocxDocument(file_path)
full_text = [para.text for para in doc.paragraphs if para.text.strip()]
content = "\n".join(full_text)
return [Document(page_content=content)]
def process_documents(document_url: str) -> PineconeVectorStore:
print("Starting data ingestion pipeline...")
temp_file_path = download_file(document_url)
if temp_file_path.endswith(".pdf"):
loader = PyPDFLoader(temp_file_path)
docs = loader.load()
elif temp_file_path.endswith(".docx"):
docs = load_docx_documents(temp_file_path)
else:
raise ValueError("Unsupported document format provided.")
print(f"Loaded {len(docs)} document pages.")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
chunked_docs = text_splitter.split_documents(docs)
print(f"Split document into {len(chunked_docs)} chunks.")
if PINECONE_INDEX_NAME not in pc.list_indexes().names():
print(f"Creating new Pinecone index: {PINECONE_INDEX_NAME}")
pc.create_index(
name=PINECONE_INDEX_NAME,
dimension=384,
metric="cosine",
spec=ServerlessSpec(cloud="aws", region="us-east-1")
)
print("Upserting document chunks to Pinecone...")
vector_store = PineconeVectorStore.from_documents(
chunked_docs, embeddings_model, index_name=PINECONE_INDEX_NAME
)
print("Document processing complete.")
os.remove(temp_file_path)
return vector_store
def create_retriever_tool(vector_store: PineconeVectorStore) -> Tool:
retriever = vector_store.as_retriever(search_kwargs={"k": 5})
def retrieve_context(query: str) -> str:
retrieved_docs = retriever.invoke(query)
return "\n\n---\n\n".join([doc.page_content for doc in retrieved_docs])
return Tool(
name="document_retriever",
func=retrieve_context,
description="Use this tool to retrieve relevant context from the policy document to answer a user's question."
) |