File size: 2,740 Bytes
3998131 |
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 |
"""
Vector database module for Module C (Letter Templates)
Stores and retrieves letter templates with embeddings.
"""
import logging
from pathlib import Path
from typing import List, Dict, Any, Optional
try:
import chromadb
CHROMADB_AVAILABLE = True
except ImportError:
CHROMADB_AVAILABLE = False
from .config import DATA_DIR
logger = logging.getLogger(__name__)
# Define Vector DB path for Module C
VECTOR_DB_DIR = DATA_DIR / "vector_db"
class TemplateVectorDB:
"""ChromaDB vector database for letter templates"""
def __init__(self, persist_directory: Path = VECTOR_DB_DIR):
"""
Initialize ChromaDB with persistent storage
"""
if not CHROMADB_AVAILABLE:
raise ImportError(
"chromadb not installed. "
"Install with: pip install chromadb"
)
self.persist_directory = Path(persist_directory)
self.persist_directory.mkdir(parents=True, exist_ok=True)
logger.info(f"Initializing Template Vector DB at {self.persist_directory}")
# Initialize ChromaDB client
self.client = chromadb.PersistentClient(
path=str(self.persist_directory)
)
# Create or get collection
self.collection_name = "nepal_letter_templates"
self.collection = self.client.get_or_create_collection(
name=self.collection_name,
metadata={"description": "Nepal letter templates for RAG generation"}
)
current_count = self.collection.count()
logger.info(f"Collection '{self.collection_name}' ready. Count: {current_count}")
def add_templates(
self,
templates: List[Dict[str, Any]],
embeddings: List[List[float]]
) -> None:
"""
Add templates with embeddings to the database
"""
if len(templates) != len(embeddings):
raise ValueError("Number of templates must match number of embeddings")
ids = [t['id'] for t in templates]
documents = [t['text'] for t in templates]
metadatas = [t['metadata'] for t in templates]
self.collection.add(
ids=ids,
documents=documents,
embeddings=embeddings,
metadatas=metadatas
)
logger.info(f"Added {len(templates)} templates to DB.")
def query_with_embedding(
self,
query_embedding: List[float],
n_results: int = 3
) -> Dict[str, Any]:
"""
Query with pre-computed embedding
"""
return self.collection.query(
query_embeddings=[query_embedding],
n_results=n_results
)
|