from pymilvus import MilvusClient, DataType
from openai import OpenAI
import fitz
import base64
import os
import uuid
class MultimodalRAG:
def __init__(self, uri: str = "./milvus.db"):
self.client = MilvusClient(uri=uri)
self.openai = OpenAI()
self.collection_name = "multimodal_rag"
self._init_collection()
def _init_collection(self):
if self.client.has_collection(self.collection_name):
return
schema = self.client.create_schema()
schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=64)
schema.add_field("content_type", DataType.VARCHAR, max_length=16)
schema.add_field("content", DataType.VARCHAR, max_length=65535)
schema.add_field("image_path", DataType.VARCHAR, max_length=512)
schema.add_field("source", DataType.VARCHAR, max_length=512)
schema.add_field("page", DataType.INT32)
schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=1536)
index_params = self.client.prepare_index_params()
index_params.add_index("embedding", index_type="AUTOINDEX", metric_type="COSINE")
self.client.create_collection(
collection_name=self.collection_name,
schema=schema,
index_params=index_params
)
def _embed(self, text: str) -> list:
"""Generate embedding using OpenAI API."""
response = self.openai.embeddings.create(
model="text-embedding-3-small",
input=[text]
)
return response.data[0].embedding
def _describe_image(self, image_path: str) -> str:
"""Generate description of image using VLM."""
with open(image_path, "rb") as f:
b64_image = base64.standard_b64encode(f.read()).decode()
response = self.openai.chat.completions.create(
model="gpt-4o",
messages=[{
"role": "user",
"content": [
{"type": "text", "text": "Describe this image in detail. If it's a chart or diagram, explain what it shows. If there's text, include it."},
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_image}"}}
]
}],
max_tokens=1000
)
return response.choices[0].message.content
def ingest_pdf(self, pdf_path: str, image_output_dir: str = "./images"):
"""Process PDF and index text + images."""
os.makedirs(image_output_dir, exist_ok=True)
doc = fitz.open(pdf_path)
source = os.path.basename(pdf_path)
data = []
for page_num, page in enumerate(doc):
text = page.get_text()
if text.strip():
chunks = self._chunk_text(text)
for chunk in chunks:
data.append({
"id": str(uuid.uuid4()),
"content_type": "text",
"content": chunk,
"image_path": "",
"source": source,
"page": page_num + 1,
"embedding": self._embed(chunk)
})
images = page.get_images()
for img_idx, img in enumerate(images):
xref = img[0]
base_image = doc.extract_image(xref)
image_bytes = base_image["image"]
image_path = f"{image_output_dir}/{source}_p{page_num+1}_img{img_idx}.png"
with open(image_path, "wb") as f:
f.write(image_bytes)
description = self._describe_image(image_path)
data.append({
"id": str(uuid.uuid4()),
"content_type": "image",
"content": description,
"image_path": image_path,
"source": source,
"page": page_num + 1,
"embedding": self._embed(description)
})
self.client.insert(collection_name=self.collection_name, data=data)
return len(data)
def _chunk_text(self, text: str, chunk_size: int = 500) -> list:
"""Split text into chunks."""
words = text.split()
chunks = []
current_chunk = []
for word in words:
current_chunk.append(word)
if len(" ".join(current_chunk)) > chunk_size:
chunks.append(" ".join(current_chunk))
current_chunk = []
if current_chunk:
chunks.append(" ".join(current_chunk))
return chunks
def retrieve(self, query: str, limit: int = 10) -> list:
"""Retrieve relevant text and images."""
embedding = self._embed(query)
results = self.client.search(
collection_name=self.collection_name,
data=[embedding],
limit=limit,
output_fields=["content_type", "content", "image_path", "source", "page"]
)
return [{
"type": hit["entity"]["content_type"],
"content": hit["entity"]["content"],
"image_path": hit["entity"]["image_path"],
"source": hit["entity"]["source"],
"page": hit["entity"]["page"],
"score": hit["distance"]
} for hit in results[0]]
def query(self, question: str, use_images: bool = True) -> dict:
"""Answer question using retrieved context."""
contexts = self.retrieve(question, limit=10)
text_contexts = [c for c in contexts if c["type"] == "text"]
image_contexts = [c for c in contexts if c["type"] == "image"]
messages = [{"role": "user", "content": []}]
context_text = "\n\n".join([
f"[{c['source']} Page {c['page']}]\n{c['content']}"
for c in text_contexts[:5]
])
messages[0]["content"].append({
"type": "text",
"text": f"Context:\n{context_text}\n"
})
if use_images and image_contexts:
for img in image_contexts[:3]:
if os.path.exists(img["image_path"]):
with open(img["image_path"], "rb") as f:
b64 = base64.standard_b64encode(f.read()).decode()
messages[0]["content"].append({
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{b64}"}
})
messages[0]["content"].append({
"type": "text",
"text": f"\nQuestion: {question}\nAnswer based on the provided context and images:"
})
response = self.openai.chat.completions.create(
model="gpt-4o",
messages=messages,
temperature=0.3
)
return {
"answer": response.choices[0].message.content,
"sources": list(set(c["source"] for c in contexts)),
"pages": list(set(c["page"] for c in contexts))
}
rag = MultimodalRAG()
rag.ingest_pdf("product_manual.pdf")
result = rag.query("What are the installation steps shown in the diagram?")
print(f"Answer: {result['answer']}")
print(f"Sources: {result['sources']}")
