File size: 6,334 Bytes
c8c00f0 | 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 | # ============================================
# DefectRAG — DefectDiffu Edition
# ============================================
# Retrieves in-context defect examples from ChromaDB and composes
# DefectDiffu text prompts (c_d) from retrieved knowledge.
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
}
# Singleton instance
_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
|