Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/multimodal/retrieval.py: 20%

1"""Cross-modal retrieval for multi-modal RAG. 

2 

3This module enables querying across modalities: 

4- Text query → Find images 

5- Image query → Find text 

6- Text query → Find audio/video 

7- Mixed modality queries 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import TYPE_CHECKING 

13 

14if TYPE_CHECKING: 

15 from lexigram.contracts.data.vector.protocols import VectorStoreProtocol 

16 

17from lexigram.ai.rag.multimodal.embeddings.multimodal import ( 

18 MultiModalEmbedder, 

19) 

20from lexigram.ai.rag.multimodal.types import ( 

21 AudioDocument, 

22 ImageDocument, 

23 Modality, 

24 MultiModalDocument, 

25 VideoDocument, 

26) 

27 

28 

29class CrossModalRetriever: 

30 """Retriever for cross-modal search. 

31 

32 Enables searching across different modalities using aligned 

33 embedding spaces (e.g., CLIP for text-image). 

34 

35 Args: 

36 embedder: MultiModalEmbedder instance 

37 vector_store: Vector store for similarity search 

38 default_top_k: Default number of results to return 

39 

40 Example: 

41 >>> retriever = CrossModalRetriever(embedder, vector_store) 

42 >>> # Text to image search 

43 >>> images = await retriever.search_by_text( 

44 ... query="a cat sitting on a mat", 

45 ... target_modality=Modality.IMAGE, 

46 ... top_k=10 

47 ... ) 

48 """ 

49 

50 def __init__( 

51 self, 

52 embedder: MultiModalEmbedder, 

53 vector_store: VectorStoreProtocol | None = None, 

54 default_top_k: int = 10, 

55 ): 

56 """Initialize cross-modal retriever.""" 

57 self.embedder = embedder 

58 self.vector_store = vector_store 

59 self.default_top_k = default_top_k 

60 

61 # Document storage (simple in-memory for now) 

62 self._documents: dict[ 

63 str, 

64 list[ImageDocument | AudioDocument | VideoDocument | MultiModalDocument], 

65 ] = { 

66 Modality.IMAGE.value: [], 

67 Modality.AUDIO.value: [], 

68 Modality.VIDEO.value: [], 

69 Modality.MULTIMODAL.value: [], 

70 } 

71 

72 # Embedding cache 

73 self._embeddings: dict[str, list[list[float]]] = { 

74 Modality.IMAGE.value: [], 

75 Modality.AUDIO.value: [], 

76 Modality.VIDEO.value: [], 

77 Modality.MULTIMODAL.value: [], 

78 } 

79 

80 async def index_image(self, image: ImageDocument) -> None: 

81 """Index an image document. 

82 

83 Args: 

84 image: ImageDocument to index 

85 """ 

86 # Embed image 

87 embedding = await self.embedder.embed_image(image) 

88 

89 # Store document and embedding 

90 self._documents[Modality.IMAGE.value].append(image) 

91 self._embeddings[Modality.IMAGE.value].append(embedding) 

92 

93 # Store embedding if available 

94 image.embedding = embedding 

95 

96 async def index_audio(self, audio: AudioDocument) -> None: 

97 """Index an audio document. 

98 

99 Args: 

100 audio: AudioDocument to index 

101 """ 

102 # Embed audio 

103 embedding = await self.embedder.embed_audio(audio) 

104 

105 # Store 

106 self._documents[Modality.AUDIO.value].append(audio) 

107 self._embeddings[Modality.AUDIO.value].append(embedding) 

108 

109 audio.embedding = embedding 

110 

111 async def index_video(self, video: VideoDocument) -> None: 

112 """Index a video document. 

113 

114 Args: 

115 video: VideoDocument to index 

116 """ 

117 # Embed video 

118 embedding = await self.embedder.embed_video(video) 

119 

120 # Store 

121 self._documents[Modality.VIDEO.value].append(video) 

122 self._embeddings[Modality.VIDEO.value].append(embedding) 

123 

124 video.embedding = embedding 

125 

126 async def index_multimodal(self, document: MultiModalDocument) -> None: 

127 """Index a multi-modal document. 

128 

129 Args: 

130 document: MultiModalDocument to index 

131 """ 

132 # Embed document 

133 embeddings = await self.embedder.embed(document) 

134 

135 # Store using fused embedding 

136 if embeddings.fused: 

137 self._documents[Modality.MULTIMODAL.value].append(document) 

138 self._embeddings[Modality.MULTIMODAL.value].append(embeddings.fused) 

139 

140 document.embeddings = embeddings 

141 

142 async def search_by_text( 

143 self, 

144 query: str, 

145 target_modality: Modality, 

146 top_k: int | None = None, 

147 ) -> list[ImageDocument | AudioDocument | VideoDocument | MultiModalDocument]: 

148 """Search for documents using text query. 

149 

150 Args: 

151 query: Text query 

152 target_modality: Modality to search in 

153 top_k: Number of results to return 

154 

155 Returns: 

156 List of documents ranked by similarity 

157 """ 

158 top_k = top_k or self.default_top_k 

159 

160 # Embed query text 

161 query_embedding = await self.embedder.embed_text(query) 

162 

163 # Search in target modality 

164 return self._similarity_search( 

165 query_embedding, 

166 target_modality, 

167 top_k, 

168 ) 

169 

170 async def search_by_image( 

171 self, 

172 image: ImageDocument, 

173 target_modality: Modality, 

174 top_k: int | None = None, 

175 ) -> list[ImageDocument | AudioDocument | VideoDocument | MultiModalDocument]: 

176 """Search for documents using image query. 

177 

178 Args: 

179 image: ImageDocument query 

180 target_modality: Modality to search in 

181 top_k: Number of results to return 

182 

183 Returns: 

184 List of documents ranked by similarity 

185 """ 

186 top_k = top_k or self.default_top_k 

187 

188 # Embed query image 

189 query_embedding = await self.embedder.embed_image(image) 

190 

191 # Search in target modality 

192 return self._similarity_search( 

193 query_embedding, 

194 target_modality, 

195 top_k, 

196 ) 

197 

198 async def search_by_embedding( 

199 self, 

200 embedding: list[float], 

201 target_modality: Modality, 

202 top_k: int | None = None, 

203 ) -> list[ImageDocument | AudioDocument | VideoDocument | MultiModalDocument]: 

204 """Search using a pre-computed embedding. 

205 

206 Args: 

207 embedding: Query embedding vector 

208 target_modality: Modality to search in 

209 top_k: Number of results to return 

210 

211 Returns: 

212 List of documents ranked by similarity 

213 """ 

214 top_k = top_k or self.default_top_k 

215 

216 return self._similarity_search( 

217 embedding, 

218 target_modality, 

219 top_k, 

220 ) 

221 

222 async def hybrid_search( 

223 self, 

224 text_query: str | None = None, 

225 image_query: ImageDocument | None = None, 

226 target_modality: Modality = Modality.MULTIMODAL, 

227 top_k: int | None = None, 

228 text_weight: float = 0.5, 

229 ) -> list[ImageDocument | AudioDocument | VideoDocument | MultiModalDocument]: 

230 """Perform hybrid search using multiple query modalities. 

231 

232 Args: 

233 text_query: Optional text query 

234 image_query: Optional image query 

235 target_modality: Modality to search in 

236 top_k: Number of results to return 

237 text_weight: Weight for text query (0-1) 

238 

239 Returns: 

240 List of documents ranked by combined similarity 

241 """ 

242 import numpy as np 

243 

244 top_k = top_k or self.default_top_k 

245 

246 # Get embeddings 

247 embeddings = [] 

248 weights = [] 

249 

250 if text_query: 

251 text_emb = await self.embedder.embed_text(text_query) 

252 embeddings.append(text_emb) 

253 weights.append(text_weight) 

254 

255 if image_query: 

256 image_emb = await self.embedder.embed_image(image_query) 

257 embeddings.append(image_emb) 

258 weights.append(1.0 - text_weight) 

259 

260 if not embeddings: 

261 return [] 

262 

263 # Combine embeddings (weighted average) 

264 np_weights = np.array(weights) / sum(weights) 

265 combined_embedding = np.average(embeddings, axis=0, weights=np_weights).tolist() 

266 

267 # Search 

268 return self._similarity_search( 

269 combined_embedding, 

270 target_modality, 

271 top_k, 

272 ) 

273 

274 def _similarity_search( 

275 self, 

276 query_embedding: list[float], 

277 target_modality: Modality, 

278 top_k: int, 

279 ) -> list[ImageDocument | AudioDocument | VideoDocument | MultiModalDocument]: 

280 """Perform similarity search in target modality. 

281 

282 Args: 

283 query_embedding: Query embedding vector 

284 target_modality: Modality to search in 

285 top_k: Number of results 

286 

287 Returns: 

288 List of top-k documents 

289 """ 

290 import numpy as np 

291 

292 modality_key = target_modality.value 

293 

294 # Get documents and embeddings 

295 documents = self._documents.get(modality_key, []) 

296 embeddings = self._embeddings.get(modality_key, []) 

297 

298 if not documents: 

299 return [] 

300 

301 # Compute similarities 

302 similarities = [] 

303 query_emb = np.array(query_embedding) 

304 

305 for doc_emb in embeddings: 

306 doc_emb_np = np.array(doc_emb) 

307 

308 # Cosine similarity 

309 similarity = np.dot(query_emb, doc_emb_np) / ( 

310 np.linalg.norm(query_emb) * np.linalg.norm(doc_emb_np) 

311 ) 

312 

313 similarities.append(similarity) 

314 

315 # Sort by similarity (descending) 

316 sorted_indices = np.argsort(similarities)[::-1] 

317 

318 # Return top-k 

319 top_k = min(top_k, len(documents)) 

320 return [documents[i] for i in sorted_indices[:top_k]] 

321 

322 def get_statistics(self) -> dict[str, int]: 

323 """Get indexing statistics. 

324 

325 Returns: 

326 Dictionary of modality -> document count 

327 """ 

328 return {modality: len(docs) for modality, docs in self._documents.items()} 

329 

330 def clear(self) -> None: 

331 """Clear all indexed documents.""" 

332 for modality in self._documents: 

333 self._documents[modality] = [] 

334 self._embeddings[modality] = []