Spaces:
Running
Running
File size: 5,247 Bytes
04dc214 | 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 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """Document storage and retrieval for RAG system."""
from typing import Optional
import uuid
import hashlib
from src.clients.gemini_embedding_client import get_document_embedding
from src.clients.qdrant_client import upsert_documents, search_similar, get_collection_info
from src.config.settings import settings
class DocumentStore:
"""Manages document storage and retrieval for RAG."""
def __init__(self, collection_name: Optional[str] = None):
"""Initialize document store.
Args:
collection_name: Qdrant collection name.
"""
self.collection_name = collection_name or settings.QDRANT_COLLECTION
self._embedding_cache = {}
def _check_document_exists(self, content: str) -> Optional[str]:
"""Check if document with same content already exists.
Args:
content: Document content to check.
Returns:
Existing document ID if found, None otherwise.
"""
content_hash = hashlib.md5(content.encode()).hexdigest()
# Check in cache first
if content_hash in self._embedding_cache:
print(f"✓ Document already exists (cached) - skipping Gemini API call")
return self._embedding_cache[content_hash]
return None
async def add_document(
self,
content: str,
title: Optional[str] = None,
source_url: Optional[str] = None,
file_path: Optional[str] = None,
section: Optional[str] = None,
tags: Optional[list] = None,
metadata: Optional[dict] = None,
) -> str:
"""Add a document to the store.
Args:
content: Document content text.
title: Document title.
source_url: URL of the source.
file_path: File path of the source.
section: Section name within the document.
tags: List of tags.
metadata: Additional metadata.
Returns:
Document ID.
"""
# Check if document already exists
existing_id = self._check_document_exists(content)
if existing_id:
return existing_id
doc_id = str(uuid.uuid4())
# Generate embedding only for new documents
print(f"→ Generating embedding for new document: {title or 'Untitled'}")
embedding = get_document_embedding(content)
# Cache the document
content_hash = hashlib.md5(content.encode()).hexdigest()
self._embedding_cache[content_hash] = doc_id
# Create payload
payload = {
"content": content,
"title": title or "",
"source_url": source_url or "",
"file_path": file_path or "",
"section": section or "",
"tags": tags or [],
"metadata": metadata or {},
"content_hash": content_hash, # Store hash for deduplication
}
# Upsert to Qdrant
upsert_documents(
collection_name=self.collection_name,
documents=[{
"id": doc_id,
"vector": embedding,
"payload": payload,
}],
)
print(f"✓ Document added successfully: {doc_id}")
return doc_id
async def add_documents_batch(
self,
documents: list[dict],
) -> list[str]:
"""Add multiple documents in batch.
Args:
documents: List of document dicts with content and metadata.
Returns:
List of document IDs.
"""
doc_ids = []
prepared_docs = []
for doc in documents:
doc_id = str(uuid.uuid4())
doc_ids.append(doc_id)
content = doc.get("content", "")
embedding = await get_embedding(content)
prepared_docs.append({
"id": doc_id,
"vector": embedding,
"payload": {
"content": content,
"title": doc.get("title", ""),
"source_url": doc.get("source_url", ""),
"file_path": doc.get("file_path", ""),
"section": doc.get("section", ""),
"tags": doc.get("tags", []),
"metadata": doc.get("metadata", {}),
},
})
await upsert_documents(
collection_name=self.collection_name,
documents=prepared_docs,
)
return doc_ids
async def search(
self,
query: str,
top_k: int = 5,
score_threshold: float = 0.4,
) -> list[dict]:
"""Search for documents similar to query.
Args:
query: Search query text.
top_k: Number of results.
score_threshold: Minimum similarity score.
Returns:
List of matching documents.
"""
query_embedding = await get_embedding(query)
return await search_similar(
collection_name=self.collection_name,
query_vector=query_embedding,
top_k=top_k,
score_threshold=score_threshold,
)
|