|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
import json |
|
|
from typing import List, Dict, Optional |
|
|
import chromadb |
|
|
from chromadb import Documents, EmbeddingFunction, Embeddings |
|
|
from sentence_transformers import SentenceTransformer |
|
|
import numpy as np |
|
|
|
|
|
|
|
|
class CLIPEmbeddingFunction(EmbeddingFunction): |
|
|
"""Custom embedding function using CLIP model.""" |
|
|
|
|
|
def __init__(self, model_name: str = "sentence-transformers/clip-ViT-B-32"): |
|
|
"""Initialize CLIP embedder.""" |
|
|
self.model = SentenceTransformer(model_name) |
|
|
|
|
|
def __call__(self, input: Documents) -> Embeddings: |
|
|
"""Generate embeddings for input documents.""" |
|
|
|
|
|
if isinstance(input, str): |
|
|
embeddings = self.model.encode([input]).tolist() |
|
|
else: |
|
|
embeddings = self.model.encode(list(input)).tolist() |
|
|
return embeddings |
|
|
|
|
|
|
|
|
class ChromaDBManager: |
|
|
"""Manage ChromaDB vector storage with persistent data.""" |
|
|
|
|
|
def __init__(self, db_dir: str = "./chroma_db"): |
|
|
"""Initialize ChromaDB with persistent storage.""" |
|
|
self.db_dir = db_dir |
|
|
os.makedirs(db_dir, exist_ok=True) |
|
|
|
|
|
|
|
|
self.client = chromadb.PersistentClient(path=db_dir) |
|
|
|
|
|
|
|
|
self.embedding_function = CLIPEmbeddingFunction( |
|
|
model_name="sentence-transformers/clip-ViT-B-32" |
|
|
) |
|
|
|
|
|
|
|
|
self.collection = self.client.get_or_create_collection( |
|
|
name="pdf_documents", |
|
|
embedding_function=self.embedding_function, |
|
|
metadata={"hnsw:space": "cosine"} |
|
|
) |
|
|
|
|
|
print(f"ChromaDB initialized. Database location: {db_dir}") |
|
|
|
|
|
def add_documents(self, documents: List[Dict]) -> None: |
|
|
"""Add documents to ChromaDB.""" |
|
|
if not documents: |
|
|
print("No documents to add") |
|
|
return |
|
|
|
|
|
doc_ids = [] |
|
|
doc_texts = [] |
|
|
doc_metadatas = [] |
|
|
|
|
|
for idx, doc in enumerate(documents): |
|
|
doc_id = f"doc_{doc.get('filename', 'unknown')}_{idx}" |
|
|
doc_text = doc.get('text', '') + " " + " ".join([table[1] for table in doc.get('tables', [])]) |
|
|
|
|
|
doc_ids.append(doc_id) |
|
|
doc_texts.append(doc_text) |
|
|
doc_metadatas.append({ |
|
|
"filename": doc.get('filename', ''), |
|
|
"page": str(doc.get('page', 0)), |
|
|
"source": "pdf" |
|
|
}) |
|
|
|
|
|
|
|
|
self.collection.add( |
|
|
ids=doc_ids, |
|
|
documents=doc_texts, |
|
|
metadatas=doc_metadatas |
|
|
) |
|
|
|
|
|
print(f"Added {len(documents)} documents to ChromaDB") |
|
|
|
|
|
def search(self, query: str, n_results: int = 5) -> List[Dict]: |
|
|
"""Search for documents similar to query.""" |
|
|
results = self.collection.query( |
|
|
query_texts=[query], |
|
|
n_results=n_results |
|
|
) |
|
|
|
|
|
retrieved_docs = [] |
|
|
if results['documents']: |
|
|
for doc, distance, metadata in zip( |
|
|
results['documents'][0], |
|
|
results['distances'][0], |
|
|
results['metadatas'][0] |
|
|
): |
|
|
retrieved_docs.append({ |
|
|
'document': doc, |
|
|
'distance': distance, |
|
|
'metadata': metadata, |
|
|
'relevance_score': 1 - distance |
|
|
}) |
|
|
|
|
|
return retrieved_docs |
|
|
|
|
|
def get_all_documents_count(self) -> int: |
|
|
"""Get total number of documents in collection.""" |
|
|
return self.collection.count() |
|
|
|
|
|
def clear_collection(self) -> None: |
|
|
"""Clear all documents from collection (for reset).""" |
|
|
self.collection.delete(where={}) |
|
|
print("Collection cleared") |
|
|
|
|
|
def get_collection_info(self) -> Dict: |
|
|
"""Get information about the collection.""" |
|
|
return { |
|
|
"name": self.collection.name, |
|
|
"document_count": self.collection.count(), |
|
|
"metadata": self.collection.metadata |
|
|
} |
|
|
|
