| |
| |
| |
| |
| |
|
|
| import chromadb |
| from sentence_transformers import SentenceTransformer |
| import json |
| from pathlib import Path |
| from typing import List, Dict, Optional |
|
|
|
|
| class DefectRAG: |
| """ |
| Retrieval-Augmented Generation for defect examples. |
| Now includes helper to build DefectDiffu defect prompts (c_d) from RAG results. |
| """ |
|
|
| def __init__( |
| self, |
| db_path: str = "data/defect_db", |
| collection: str = "defect_patches", |
| model: str = "all-MiniLM-L6-v2" |
| ): |
| self.db_path = db_path |
| self.collection_name = collection |
| self._client = None |
| self._collection = None |
| self._encoder = None |
|
|
| @property |
| def client(self): |
| if self._client is None: |
| self._client = chromadb.PersistentClient(path=self.db_path) |
| return self._client |
|
|
| @property |
| def collection(self): |
| if self._collection is None: |
| self._collection = self.client.get_collection(self.collection_name) |
| return self._collection |
|
|
| @property |
| def encoder(self): |
| if self._encoder is None: |
| self._encoder = SentenceTransformer('all-MiniLM-L6-v2') |
| return self._encoder |
|
|
| def _build_where_filter( |
| self, |
| commercial_only: bool = True, |
| domain_filter: Optional[str] = None |
| ) -> Optional[Dict]: |
| conditions = [] |
| if commercial_only: |
| conditions.append({"commercial_ok": True}) |
| if domain_filter and domain_filter != "general": |
| conditions.append({"domain": domain_filter}) |
| if len(conditions) == 0: |
| return None |
| elif len(conditions) == 1: |
| return conditions[0] |
| else: |
| return {"$and": conditions} |
|
|
| def retrieve( |
| self, |
| defect_plan, |
| k: int = 3, |
| domain_filter: Optional[str] = None, |
| commercial_only: bool = True |
| ) -> List[Dict]: |
| """Retrieve top-k matching defect examples.""" |
| query_parts = [ |
| getattr(defect_plan, 'artifact_type', defect_plan.defect_type), |
| defect_plan.description, |
| "on", |
| defect_plan.target_entity |
| ] |
| query = " ".join(query_parts) |
| query_emb = self.encoder.encode(query) |
| where_filter = self._build_where_filter(commercial_only, domain_filter) |
|
|
| results = self.collection.query( |
| query_embeddings=[query_emb.tolist()], |
| n_results=k, |
| where=where_filter |
| ) |
|
|
| examples = [] |
| for i, meta in enumerate(results['metadatas'][0]): |
| paths = json.loads(meta['paths']) if isinstance(meta.get('paths'), str) else meta.get('paths', {}) |
| examples.append({ |
| 'paths': paths, |
| 'caption': meta.get('caption', ''), |
| 'domain': meta.get('domain', 'unknown'), |
| 'license': meta.get('license', 'unknown'), |
| 'source': meta.get('source', 'unknown'), |
| 'defect_name': meta.get('defect_name', 'unknown'), |
| 'score': results['distances'][0][i] if results.get('distances') else None, |
| 'metadata': {k: v for k, v in meta.items() if k not in {'paths', 'caption', 'domain', 'license', 'source', 'defect_name'}} |
| }) |
| return examples |
|
|
| def retrieve_by_text( |
| self, |
| text: str, |
| k: int = 3, |
| domain_filter: Optional[str] = None |
| ) -> List[Dict]: |
| """Direct text search (for debugging/testing).""" |
| query_emb = self.encoder.encode(text) |
| where_filter = self._build_where_filter(True, domain_filter) |
| results = self.collection.query( |
| query_embeddings=[query_emb.tolist()], |
| n_results=k, |
| where=where_filter |
| ) |
| examples = [] |
| for i, meta in enumerate(results['metadatas'][0]): |
| paths = json.loads(meta['paths']) if isinstance(meta.get('paths'), str) else meta.get('paths', {}) |
| examples.append({ |
| 'paths': paths, |
| 'caption': meta.get('caption', ''), |
| 'domain': meta.get('domain', 'unknown'), |
| 'score': results['distances'][0][i] if results.get('distances') else None |
| }) |
| return examples |
|
|
| def compose_defect_prompt( |
| self, |
| base_description: str, |
| examples: List[Dict], |
| max_examples: int = 1 |
| ) -> str: |
| """ |
| Build a DefectDiffu defect prompt (c_d) by enriching the base description |
| with captions from retrieved RAG examples. |
| |
| Example output: |
| "A photo of a small transparent bubble trapped under glass, similar to |
| a spherical air pocket with dark meniscus ring" |
| """ |
| if not examples: |
| return f"A photo of {base_description}" |
|
|
| captions = [ex.get('caption', '') for ex in examples[:max_examples] if ex.get('caption')] |
| if captions: |
| enriched = f"{base_description}, similar to {captions[0]}" |
| return f"A photo of {enriched}" |
| return f"A photo of {base_description}" |
|
|
| def get_stats(self) -> Dict: |
| """Get DB statistics.""" |
| count = self.collection.count() |
| results = self.collection.get() |
| domains = {} |
| licenses = {} |
| sources = {} |
| for meta in results["metadatas"]: |
| domains[meta.get("domain", "unknown")] = domains.get(meta.get("domain"), 0) + 1 |
| licenses[meta.get("license", "unknown")] = licenses.get(meta.get("license"), 0) + 1 |
| sources[meta.get("source", "unknown")] = sources.get(meta.get("source"), 0) + 1 |
| return { |
| "total_entries": count, |
| "domains": domains, |
| "licenses": licenses, |
| "sources": sources |
| } |
|
|
|
|
| |
| _rag_instance = None |
|
|
| def get_rag(db_path="data/defect_db") -> DefectRAG: |
| """Get or create singleton RAG instance.""" |
| global _rag_instance |
| if _rag_instance is None: |
| _rag_instance = DefectRAG(db_path=db_path) |
| return _rag_instance |
|
|
