File size: 13,992 Bytes
8275526
 
 
 
 
289c45a
8275526
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ca362fb
 
 
 
 
 
 
 
 
 
8275526
289c45a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8275526
289c45a
 
 
 
 
 
 
8275526
289c45a
 
 
 
 
 
8275526
 
 
 
 
 
 
 
 
 
ca362fb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8275526
 
 
 
 
ca362fb
8275526
 
 
 
 
 
 
 
ca362fb
 
 
 
 
8275526
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ca362fb
 
8275526
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ca362fb
 
8275526
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ca362fb
8275526
 
 
 
 
 
 
ca362fb
8275526
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ca362fb
8275526
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
import logging
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings
import uuid
import os
from config import EMBEDDING_MODEL, CHROMA_PERSIST_DIRECTORY, COLLECTION_NAME

# Cấu hình logging
logger = logging.getLogger(__name__)

# Global instance để implement singleton pattern
_embedding_model_instance = None

def get_embedding_model():
    """
    Singleton pattern để đảm bảo chỉ có một instance của EmbeddingModel
    """
    global _embedding_model_instance
    if _embedding_model_instance is None:
        logger.info("Khởi tạo EmbeddingModel instance lần đầu")
        _embedding_model_instance = EmbeddingModel()
    else:
        logger.debug("Sử dụng EmbeddingModel instance đã có")
    return _embedding_model_instance

class EmbeddingModel:
    def __init__(self):
        """Khởi tạo embedding model và ChromaDB client"""
        logger.info(f"Đang khởi tạo embedding model: {EMBEDDING_MODEL}")
        
        try:
            # Khởi tạo sentence transformer với trust_remote_code=True
            self.model = SentenceTransformer(EMBEDDING_MODEL, trust_remote_code=True)
            logger.info("Đã tải sentence transformer model")
        except Exception as e:
            logger.error(f"Lỗi khởi tạo model: {e}")
            # Thử với cache_folder explicit
            cache_dir = os.getenv('SENTENCE_TRANSFORMERS_HOME', '/app/.cache/sentence-transformers')
            self.model = SentenceTransformer(EMBEDDING_MODEL, cache_folder=cache_dir, trust_remote_code=True)
            logger.info("Đã tải sentence transformer model với cache folder explicit")
        
        # Đảm bảo thư mục ChromaDB tồn tại và có quyền ghi
        try:
            os.makedirs(CHROMA_PERSIST_DIRECTORY, exist_ok=True)
            # Test ghi file để kiểm tra permission
            test_file = os.path.join(CHROMA_PERSIST_DIRECTORY, 'test_permission.tmp')
            with open(test_file, 'w') as f:
                f.write('test')
            os.remove(test_file)
            logger.info(f"Thư mục ChromaDB đã sẵn sàng: {CHROMA_PERSIST_DIRECTORY}")
        except Exception as e:
            logger.error(f"Lỗi tạo/kiểm tra thư mục ChromaDB: {e}")
            # Fallback to /tmp directory
            import tempfile
            CHROMA_PERSIST_DIRECTORY = os.path.join(tempfile.gettempdir(), 'chroma_db')
            os.makedirs(CHROMA_PERSIST_DIRECTORY, exist_ok=True)
            logger.warning(f"Sử dụng thư mục tạm thời: {CHROMA_PERSIST_DIRECTORY}")
        
        # Khởi tạo ChromaDB client với persistent storage
        try:
            self.chroma_client = chromadb.PersistentClient(
                path=CHROMA_PERSIST_DIRECTORY,
                settings=Settings(
                    anonymized_telemetry=False,
                    allow_reset=True
                )
            )
            logger.info(f"Đã kết nối ChromaDB tại: {CHROMA_PERSIST_DIRECTORY}")
        except Exception as e:
            logger.error(f"Lỗi kết nối ChromaDB: {e}")
            # Fallback to in-memory client
            logger.warning("Fallback to in-memory ChromaDB client")
            self.chroma_client = chromadb.Client()
        
        # Lấy hoặc tạo collection
        try:
            self.collection = self.chroma_client.get_collection(name=COLLECTION_NAME)
            logger.info(f"Đã kết nối collection '{COLLECTION_NAME}' với {self.collection.count()} items")
        except Exception:
            logger.warning(f"Collection '{COLLECTION_NAME}' không tồn tại, tạo mới...")
            self.collection = self.chroma_client.create_collection(name=COLLECTION_NAME)
            logger.info(f"Đã tạo collection mới: {COLLECTION_NAME}")
    
    def _add_prefix_to_text(self, text, is_query=True):
        """
        Thêm prefix cho text theo yêu cầu của multilingual-e5-base
        """
        # Kiểm tra xem text đã có prefix chưa
        if text.startswith(('query:', 'passage:')):
            return text
        
        # Thêm prefix phù hợp
        if is_query:
            return f"query: {text}"
        else:
            return f"passage: {text}"
    
    def encode(self, texts, is_query=True):
        """
        Encode văn bản thành embeddings
        
        Args:
            texts (str or list): Văn bản hoặc danh sách văn bản cần encode
            is_query (bool): True nếu là query, False nếu là passage
            
        Returns:
            list: Embeddings vector
        """
        try:
            if isinstance(texts, str):
                texts = [texts]
            
            # Thêm prefix cho texts
            processed_texts = [self._add_prefix_to_text(text, is_query) for text in texts]
            
            logger.debug(f"Đang encode {len(processed_texts)} văn bản")
            embeddings = self.model.encode(processed_texts, show_progress_bar=False, normalize_embeddings=True)
            
            return embeddings.tolist()
            
        except Exception as e:
            logger.error(f"Lỗi encode văn bản: {e}")
            raise
    
    def search(self, query, top_k=5, age_filter=None):
        """
        Tìm kiếm văn bản tương tự trong ChromaDB
        
        Args:
            query (str): Câu hỏi cần tìm kiếm
            top_k (int): Số lượng kết quả trả về
            age_filter (int): Lọc theo độ tuổi (optional)
            
        Returns:
            list: Danh sách kết quả tìm kiếm
        """
        try:
            logger.debug(f"Dang tim kiem cho query: {query[:50]}...")
            
            # Encode query thành embedding (với prefix query:)
            query_embedding = self.encode(query, is_query=True)[0]
            
            # Tạo where clause cho age filter
            where_clause = None
            if age_filter:
                where_clause = {
                    "$and": [
                        {"age_min": {"$lte": age_filter}},
                        {"age_max": {"$gte": age_filter}}
                    ]
                }
            
            # Thực hiện search trong ChromaDB
            search_results = self.collection.query(
                query_embeddings=[query_embedding],
                n_results=top_k,
                where=where_clause,
                include=['documents', 'metadatas', 'distances']
            )
            
            if not search_results or not search_results['documents']:
                logger.warning("Khong tim thay ket qua nao")
                return []
            
            # Format kết quả
            results = []
            documents = search_results['documents'][0]
            metadatas = search_results['metadatas'][0]
            distances = search_results['distances'][0]
            
            for i, (doc, metadata, distance) in enumerate(zip(documents, metadatas, distances)):
                results.append({
                    'document': doc,
                    'metadata': metadata or {},
                    'distance': distance,
                    'similarity': 1 - distance,  # Chuyển distance thành similarity
                    'rank': i + 1
                })
            
            logger.info(f"Tim thay {len(results)} ket qua cho query")
            return results
            
        except Exception as e:
            logger.error(f"Loi tim kiem: {e}")
            return []
    
    def add_documents(self, documents, metadatas=None, ids=None):
        """
        Thêm documents vào ChromaDB
        
        Args:
            documents (list): Danh sách văn bản
            metadatas (list): Danh sách metadata tương ứng
            ids (list): Danh sách ID tương ứng (optional)
            
        Returns:
            bool: True nếu thành công
        """
        try:
            if not documents:
                logger.warning("Không có documents để thêm")
                return False
            
            # Tạo IDs nếu không được cung cấp
            if not ids:
                ids = [str(uuid.uuid4()) for _ in documents]
            
            # Tạo metadatas rỗng nếu không được cung cấp
            if not metadatas:
                metadatas = [{} for _ in documents]
            
            logger.info(f"Đang thêm {len(documents)} documents vào ChromaDB")
            
            # Encode documents thành embeddings (với prefix passage:)
            embeddings = self.encode(documents, is_query=False)
            
            # Thêm vào collection
            self.collection.add(
                embeddings=embeddings,
                documents=documents,
                metadatas=metadatas,
                ids=ids
            )
            
            logger.info(f"Đã thêm thành công {len(documents)} documents")
            return True
            
        except Exception as e:
            logger.error(f"Lỗi thêm documents: {e}")
            return False
    
    def index_chunks(self, chunks):
        """
        Index các chunks dữ liệu vào ChromaDB
        """
        try:
            if not chunks:
                logger.warning("Không có chunks để index")
                return False
            
            documents = []
            metadatas = []
            ids = []
            
            for chunk in chunks:
                if not chunk.get('content'):
                    logger.warning(f"Chunk thiếu content: {chunk}")
                    continue
                
                documents.append(chunk['content'])
                
                # Lấy metadata đã được chuẩn bị sẵn
                metadata = chunk.get('metadata', {})
                metadatas.append(metadata)
                
                # Sử dụng ID có sẵn hoặc tạo mới
                chunk_id = chunk.get('id') or str(uuid.uuid4())
                ids.append(chunk_id)
            
            if not documents:
                logger.warning("Không có documents hợp lệ để index")
                return False
            
            # Batch processing để tránh overload
            batch_size = 100
            total_batches = (len(documents) + batch_size - 1) // batch_size
            
            for i in range(0, len(documents), batch_size):
                batch_docs = documents[i:i + batch_size]
                batch_metas = metadatas[i:i + batch_size]
                batch_ids = ids[i:i + batch_size]
                
                batch_num = (i // batch_size) + 1
                logger.info(f"Đang xử lý batch {batch_num}/{total_batches} ({len(batch_docs)} items)")
                
                success = self.add_documents(batch_docs, batch_metas, batch_ids)
                if not success:
                    logger.error(f"Lỗi xử lý batch {batch_num}")
                    return False
            
            logger.info(f"Đã index thành công {len(documents)} chunks")
            return True
            
        except Exception as e:
            logger.error(f"Lỗi index chunks: {e}")
            return False
    
    def count(self):
        """Đếm số lượng documents trong collection"""
        try:
            return self.collection.count()
        except Exception as e:
            logger.error(f"Lỗi đếm documents: {e}")
            return 0
    
    def delete_collection(self):
        """Xóa collection hiện tại"""
        try:
            logger.warning(f"Đang xóa collection: {COLLECTION_NAME}")
            self.chroma_client.delete_collection(name=COLLECTION_NAME)
            
            # Tạo lại collection mới
            self.collection = self.chroma_client.create_collection(name=COLLECTION_NAME)
            logger.info("Đã tạo lại collection mới")
            
            return True
            
        except Exception as e:
            logger.error(f"Lỗi xóa collection: {e}")
            return False
    
    def get_stats(self):
        """Lấy thống kê về collection"""
        try:
            total_count = self.count()
            
            # Lấy sample để phân tích metadata
            sample_results = self.collection.get(limit=min(100, total_count))
            
            # Thống kê content types
            content_types = {}
            chapters = {}
            age_groups = {}
            
            if sample_results and sample_results.get('metadatas'):
                for metadata in sample_results['metadatas']:
                    if not metadata:
                        continue
                    
                    # Content type stats
                    content_type = metadata.get('content_type', 'unknown')
                    content_types[content_type] = content_types.get(content_type, 0) + 1
                    
                    # Chapter stats
                    chapter = metadata.get('chapter', 'unknown')
                    chapters[chapter] = chapters.get(chapter, 0) + 1
                    
                    # Age group stats
                    age_group = metadata.get('age_group', 'unknown')
                    age_groups[age_group] = age_groups.get(age_group, 0) + 1
            
            return {
                'total_documents': total_count,
                'content_types': content_types,
                'chapters': chapters,
                'age_groups': age_groups,
                'collection_name': COLLECTION_NAME,
                'embedding_model': EMBEDDING_MODEL
            }
            
        except Exception as e:
            logger.error(f"Lỗi lấy stats: {e}")
            return {
                'total_documents': 0,
                'error': str(e)
            }