@samitouri / QOSami_Agents_Zero-1 / commits / 9c8703cd

rag tool finalizing

frdel committed Jun 15, 2025 at 22:00 UTC 9c8703cd7c7f97221e7420863e0dcc32eeac0ed5

4 files changed +255 -316

python/api/message.py

+1 -1

@@ -31,7 +31,7 @@ class Message(ApiHandler):

31 attachment_paths = []

33 upload_folder_int = "/a0/tmp/uploads"

34 - upload_folder_ext = files.get_abs_path("tmp/uploads")

34 + upload_folder_ext = files.get_abs_path("tmp/uploads") # for development environment

36 if attachments:

37 os.makedirs(upload_folder_ext, exist_ok=True)

python/helpers/document_query.py

+193 -298

@@ -4,6 +4,8 @@ import asyncio

4 import aiohttp

5 import json

7 +from python.helpers.vector_db import VectorDB

8 +

9 os.environ["USER_AGENT"] = "@mixedbread-ai/unstructured" # noqa E402

10 from langchain_unstructured import UnstructuredLoader # noqa E402

@@ -32,15 +34,14 @@ from langchain_community.vectorstores.utils import (

34 from langchain_core.embeddings import Embeddings

36 from python.helpers.print_style import PrintStyle

35 -from python.helpers import files #, rfc_files

36 -rfc_files = files # TODO: fix

37 -

37 +from python.helpers import files

38 from agent import Agent

39 -import models

40 from langchain.text_splitter import RecursiveCharacterTextSplitter

43 +DEFAULT_SEARCH_THRESHOLD = 0.6

44 +

45 class DocumentQueryStore:

46 """

47 FAISS Store for document query results.

@@ -55,107 +56,25 @@ class DocumentQueryStore:

56 _stores: dict[str, "DocumentQueryStore"] = {}

58 @staticmethod

58 - async def get(agent: Agent):

59 - """Get or create a DocumentQueryStore instance for the specified agent."""

59 + def get(agent: Agent):

60 + """Create a DocumentQueryStore instance for the specified agent."""

61 if not agent or not agent.config:

62 raise ValueError("Agent and agent config must be provided")

63 - memory_subdir = agent.config.memory_subdir or "default"

64 - store_key = f"{memory_subdir}/document_query"

65 -

66 - if store_key not in DocumentQueryStore._stores:

67 - # Initialize embeddings model from agent config

68 - embeddings_model = agent.get_embedding_model()

69 -

70 - # Initialize store

71 - store = DocumentQueryStore(agent, embeddings_model, memory_subdir)

72 - DocumentQueryStore._stores[store_key] = store

73 - return store

74 - else:

75 - return DocumentQueryStore._stores[store_key]

76 -

77 - @staticmethod

78 - async def reload(agent: Agent):

79 - """Reload the DocumentQueryStore for the specified agent."""

80 - memory_subdir = agent.config.memory_subdir or "default"

81 - store_key = f"{memory_subdir}/document_query"

82 -

83 - if store_key in DocumentQueryStore._stores:

84 - del DocumentQueryStore._stores[store_key]

85 -

86 - return await DocumentQueryStore.get(agent)

64 + # Initialize store

65 + store = DocumentQueryStore(agent)

66 + return store

68 def __init__(

69 self,

70 agent: Agent,

91 - embeddings_model: Embeddings,

92 - memory_subdir: str,

71 ):

72 """Initialize a DocumentQueryStore instance."""

73 self.agent = agent

96 - self.memory_subdir = memory_subdir

97 -

98 - # Get directory paths

99 - db_dir = self._get_db_dir()

100 - em_dir = os.path.join(db_dir, "embeddings")

101 -

102 - # Create directories

103 - os.makedirs(db_dir, exist_ok=True)

104 - os.makedirs(em_dir, exist_ok=True)

105 -

106 - # Setup embeddings cache

107 - store = LocalFileStore(em_dir)

108 - self.embeddings = CacheBackedEmbeddings.from_bytes_store(

109 - embeddings_model,

110 - store,

111 - namespace=f"document_query_{getattr(embeddings_model, 'model', getattr(embeddings_model, 'model_name', 'default'))}"

112 - )

113 -

114 - # Initialize vector store

115 - index_path = os.path.join(db_dir, "index.faiss")

116 - docstore_path = os.path.join(db_dir, "docstore.json")

117 -

118 - if os.path.exists(index_path) and os.path.exists(docstore_path):

119 - PrintStyle.standard(f"Loading existing vector store from {db_dir}")

120 - try:

121 - self.vectorstore = FAISS.load_local(

122 - folder_path=db_dir,

123 - embeddings=self.embeddings,

124 - allow_dangerous_deserialization=True,

125 - distance_strategy=DistanceStrategy.COSINE,

126 - )

127 - except Exception as e:

128 - PrintStyle.error(f"Error loading vector store: {str(e)}")

129 - self._initialize_new_vectorstore()

130 - else:

131 - PrintStyle.standard(f"Creating new vector store in '{db_dir}'")

132 - self._initialize_new_vectorstore()

133 -

134 - def _initialize_new_vectorstore(self):

135 - """Initialize a new vector store."""

136 - dimension = len(self.embeddings.embed_query("test"))

137 - index = faiss.IndexFlatIP(dimension)

138 - self.vectorstore = FAISS(

139 - embedding_function=self.embeddings,

140 - index=index,

141 - docstore=InMemoryDocstore(),

142 - index_to_docstore_id={},

143 - distance_strategy=DistanceStrategy.COSINE,

144 - )

145 -

146 - def _get_db_dir(self) -> str:

147 - """Get the absolute path to the database directory."""

148 - return files.get_abs_path("memory", self.memory_subdir, "document_query")

149 -

150 - def _save_vectorstore(self):

151 - """Save the vector store to disk."""

152 - db_dir = self._get_db_dir()

153 - PrintStyle.standard(f"Saving vector store to {db_dir}")

154 - self.vectorstore.save_local(folder_path=db_dir)

155 - PrintStyle.standard(f"Vector store saved with {len(self.vectorstore.index_to_docstore_id)} documents")

74 + self.vector_db: VectorDB | None = None

76 @staticmethod

158 - def _normalize_uri(uri: str) -> str:

77 + def normalize_uri(uri: str) -> str:

78 """

79 Normalize a document URI to ensure consistent lookup.

@@ -166,7 +85,7 @@ class DocumentQueryStore:

85 Normalized URI

86 """

87 # Convert to lowercase

169 - normalized = uri.lower()

88 + normalized = uri.strip() # uri.lower()

90 # Parse the URL to get scheme

91 parsed = urlparse(normalized)

@@ -174,17 +93,23 @@ class DocumentQueryStore:

94 # Normalize based on scheme

95 if scheme == "file":

177 - if not normalized.startswith("file:"):

178 - normalized = "file:" + normalized

179 - if normalized.startswith("file://"):

180 - normalized = normalized.replace("file://", "file:")

96 + path = files.fix_dev_path(

97 + normalized.removeprefix("file://").removeprefix("file:")

98 + )

99 + normalized = f"file://{path}"

100 +

101 elif scheme in ["http", "https"]:

102 # Always use https for web URLs

183 - normalized = normalized.replace("http://", "https://")

103 + normalized = normalized.replace("http://", "https://") # TODO why?

104

105 return normalized

106

187 - async def add_document(self, text: str, document_uri: str, metadata: dict | None = None) -> bool:

107 + def init_vector_db(self):

108 + return VectorDB(self.agent)

109 +

110 + async def add_document(

111 + self, text: str, document_uri: str, metadata: dict | None = None

112 + ) -> bool:

113 """

114 Add a document to the store with the given URI.

115

@@ -197,7 +122,7 @@ class DocumentQueryStore:

122 True if successful, False otherwise

123 """

124 # Normalize the URI

200 - document_uri = self._normalize_uri(document_uri)

125 + document_uri = self.normalize_uri(document_uri)

126

127 # Delete existing document if it exists to avoid duplicates

128 await self.delete_document(document_uri)

@@ -209,8 +134,7 @@ class DocumentQueryStore:

134

135 # Split text into chunks

136 text_splitter = RecursiveCharacterTextSplitter(

212 - chunk_size=self.DEFAULT_CHUNK_SIZE,

213 - chunk_overlap=self.DEFAULT_CHUNK_OVERLAP

137 + chunk_size=self.DEFAULT_CHUNK_SIZE, chunk_overlap=self.DEFAULT_CHUNK_OVERLAP

138 )

139 chunks = text_splitter.split_text(text)

140

@@ -230,14 +154,17 @@ class DocumentQueryStore:

154 try:

155 docs_text = "".join(chunk.page_content for chunk in docs)

156 await self.agent.rate_limiter(

233 - model_config=self.agent.config.embeddings_model,

234 - input=docs_text

157 + model_config=self.agent.config.embeddings_model, input=docs_text

158 )

159

237 - # Add documents to vector store

238 - self.vectorstore.add_documents(documents=docs)

239 - self._save_vectorstore()

240 - PrintStyle.standard(f"Added document '{document_uri}' with {len(docs)} chunks")

160 + # Initialize vector db if not already initialized

161 + if not self.vector_db:

162 + self.vector_db = self.init_vector_db()

163 +

164 + ids = await self.vector_db.insert_documents(docs)

165 + PrintStyle.standard(

166 + f"Added document '{document_uri}' with {len(docs)} chunks"

167 + )

168 return True

169 except Exception as e:

170 PrintStyle.error(f"Error adding document '{document_uri}': {str(e)}")

@@ -253,8 +180,13 @@ class DocumentQueryStore:

180 Returns:

181 The complete document if found, None otherwise

182 """

183 +

184 + # DB not initialized, no documents inside

185 + if not self.vector_db:

186 + return None

187 +

188 # Normalize the URI

257 - document_uri = self._normalize_uri(document_uri)

189 + document_uri = self.normalize_uri(document_uri)

190

191 # Get all chunks for this document

192 docs = await self._get_document_chunks(document_uri)

@@ -283,14 +215,19 @@ class DocumentQueryStore:

215 Returns:

216 List of document chunks

217 """

218 +

219 + # DB not initialized, no documents inside

220 + if not self.vector_db:

221 + return []

222 +

223 # Normalize the URI

287 - document_uri = self._normalize_uri(document_uri)

224 + document_uri = self.normalize_uri(document_uri)

225

289 - # Access docstore directly

290 - chunks = []

291 - for doc_id, doc in self.vectorstore.docstore._dict.items(): # type: ignore

292 - if isinstance(doc.metadata, dict) and doc.metadata.get("document_uri") == document_uri:

293 - chunks.append(doc)

226 + # get docs from vector db

227 +

228 + chunks = await self.vector_db.search_by_metadata(

229 + filter=f"document_uri == '{document_uri}'",

230 + )

231

232 PrintStyle.standard(f"Found {len(chunks)} chunks for document: {document_uri}")

233 return chunks

@@ -305,8 +242,13 @@ class DocumentQueryStore:

242 Returns:

243 True if the document exists, False otherwise

244 """

245 +

246 + # DB not initialized, no documents inside

247 + if not self.vector_db:

248 + return False

249 +

250 # Normalize the URI

309 - document_uri = self._normalize_uri(document_uri)

251 + document_uri = self.normalize_uri(document_uri)

252

253 chunks = await self._get_document_chunks(document_uri)

254 return len(chunks) > 0

@@ -321,84 +263,36 @@ class DocumentQueryStore:

263 Returns:

264 True if deleted, False if not found

265 """

266 +

267 + # DB not initialized, no documents inside

268 + if not self.vector_db:

269 + return False

270 +

271 # Normalize the URI

325 - document_uri = self._normalize_uri(document_uri)

272 + document_uri = self.normalize_uri(document_uri)

273

327 - chunks = await self._get_document_chunks(document_uri)

274 + chunks = await self.vector_db.search_by_metadata(

275 + filter=f"document_uri == '{document_uri}'",

276 + )

277 if not chunks:

278 return False

279

280 # Collect IDs to delete

332 - ids_to_delete = []

333 - for chunk in chunks:

334 - for doc_id, doc_ref in self.vectorstore.docstore._dict.items(): # type: ignore

335 - if doc_ref == chunk:

336 - ids_to_delete.append(doc_id)

281 + ids_to_delete = [chunk.metadata["id"] for chunk in chunks]

282

283 # Delete from vector store

284 if ids_to_delete:

340 - self.vectorstore.delete(ids_to_delete)

341 - self._save_vectorstore()

342 - PrintStyle.standard(f"Deleted document '{document_uri}' with {len(ids_to_delete)} chunks")

285 + dels = await self.vector_db.delete_documents_by_ids(ids_to_delete)

286 + PrintStyle.standard(

287 + f"Deleted document '{document_uri}' with {len(dels)} chunks"

288 + )

289 return True

290

291 return False

292

347 - async def expire_documents(self, older_than_days: float) -> int:

348 - """

349 - Delete documents older than the specified number of days.

350 -

351 - Args:

352 - older_than_days: Number of days (can be fractional) before current time

353 -

354 - Returns:

355 - Number of documents deleted

356 - """

357 - if older_than_days <= 0:

358 - return 0

359 -

360 - # Calculate cutoff timestamp

361 - cutoff_date = datetime.now().timestamp() - (older_than_days * 24 * 60 * 60)

362 -

363 - # Find expired documents

364 - expired_uris = set()

365 -

366 - # Check all documents in the store

367 - for doc_id, doc in self.vectorstore.docstore._dict.items(): # type: ignore

368 - if not isinstance(doc.metadata, dict):

369 - continue

370 -

371 - # Only process each document once (first chunk)

372 - if doc.metadata.get("chunk_index", 0) != 0:

373 - continue

374 -

375 - doc_uri = doc.metadata.get("document_uri")

376 - if not doc_uri:

377 - continue

378 -

379 - try:

380 - # Check timestamp

381 - timestamp_str = doc.metadata.get("timestamp")

382 - if not timestamp_str:

383 - continue

384 -

385 - doc_timestamp = datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S").timestamp()

386 - if doc_timestamp < cutoff_date:

387 - expired_uris.add(doc_uri)

388 - except (ValueError, TypeError):

389 - # Skip documents with invalid timestamps

390 - continue

391 -

392 - # Delete expired documents

393 - deleted_count = 0

394 - for uri in expired_uris:

395 - if await self.delete_document(uri):

396 - deleted_count += 1

397 -

398 - PrintStyle.standard(f"Expired {deleted_count} documents older than {older_than_days} days")

399 - return deleted_count

400 -

401 - async def search_documents(self, query: str, limit: int = 10, threshold: float = 0.5) -> List[Document]:

293 + async def search_documents(

294 + self, query: str, limit: int = 10, threshold: float = 0.5, filter: str = ""

295 + ) -> List[Document]:

296 """

297 Search for documents similar to the query across the entire store.

298

@@ -410,34 +304,30 @@ class DocumentQueryStore:

304 Returns:

305 List of matching documents

306 """

307 +

308 + # DB not initialized, no documents inside

309 + if not self.vector_db:

310 + return []

311 +

312 # Handle empty query

313 if not query:

415 - PrintStyle.standard("Empty search query, returning empty results")

314 return []

315

418 - # Apply rate limiter

419 - await self.agent.rate_limiter(

420 - model_config=self.agent.config.embeddings_model,

421 - input=query

422 - )

423 -

316 # Perform search

317 try:

426 - results = self.vectorstore.similarity_search_with_score(

427 - query=query,

428 - k=limit,

429 - score_threshold=threshold

318 + results = await self.vector_db.search_by_similarity_threshold(

319 + query=query, limit=limit, threshold=threshold, filter=filter

320 )

321

432 - # Extract documents from results (which are (doc, score) pairs)

433 - docs = [doc for doc, score in results]

434 - PrintStyle.standard(f"Search '{query}' returned {len(docs)} results")

435 - return docs

322 + PrintStyle.standard(f"Search '{query}' returned {len(results)} results")

323 + return results

324 except Exception as e:

325 PrintStyle.error(f"Error searching documents: {str(e)}")

326 return []

327

440 - async def search_document(self, document_uri: str, query: str, limit: int = 10, threshold: float = 0.5) -> List[Document]:

328 + async def search_document(

329 + self, document_uri: str, query: str, limit: int = 10, threshold: float = 0.5

330 + ) -> List[Document]:

331 """

332 Search for content within a specific document.

333

@@ -450,59 +340,10 @@ class DocumentQueryStore:

340 Returns:

341 List of matching document chunks

342 """

453 - # Normalize the URI

454 - document_uri = self._normalize_uri(document_uri)

455 -

456 - # Handle empty query

457 - if not query:

458 - PrintStyle.standard("Empty search query, returning empty results")

459 - return []

460 -

461 - # Check if document exists

462 - if not await self.document_exists(document_uri):

463 - PrintStyle.error(f"Document not found: {document_uri}")

464 - return []

465 -

466 - # Apply rate limiter

467 - await self.agent.rate_limiter(

468 - model_config=self.agent.config.embeddings_model,

469 - input=query

343 + return await self.search_documents(

344 + query, limit, threshold, f"document_uri == '{document_uri}'"

345 )

346

472 - # Perform search with document filter

473 - try:

474 - # Create metadata filter function

475 - def filter_fn(doc_metadata):

476 - return doc_metadata.get("document_uri") == document_uri

477 -

478 - results = self.vectorstore.similarity_search_with_score(

479 - query=query,

480 - k=limit,

481 - score_threshold=threshold,

482 - filter=filter_fn

483 - )

484 -

485 - # Extract documents from results

486 - docs = [doc for doc, score in results]

487 - PrintStyle.standard(f"Search '{query}' in document '{document_uri}' returned {len(docs)} results")

488 -

489 - # Try with lower threshold if no results

490 - if not docs and threshold > 0.3:

491 - PrintStyle.standard("No results found, trying with lower threshold (0.3)")

492 - results = self.vectorstore.similarity_search_with_score(

493 - query=query,

494 - k=limit,

495 - score_threshold=0.3,

496 - filter=filter_fn

497 - )

498 - docs = [doc for doc, score in results]

499 - PrintStyle.standard(f"Retry search returned {len(docs)} results")

500 -

501 - return docs

502 - except Exception as e:

503 - PrintStyle.error(f"Error searching within document: {str(e)}")

504 - return []

505 -

347 async def list_documents(self) -> List[str]:

348 """

349 Get a list of all document URIs in the store.

@@ -510,9 +351,13 @@ class DocumentQueryStore:

351 Returns:

352 List of document URIs

353 """

354 + # DB not initialized, no documents inside

355 + if not self.vector_db:

356 + return []

357 +

358 # Extract unique URIs

359 uris = set()

515 - for doc in self.vectorstore.docstore._dict.values(): # type: ignore

360 + for doc in self.vector_db.db.get_all_docs().values():

361 if isinstance(doc.metadata, dict):

362 uri = doc.metadata.get("document_uri")

363 if uri:

@@ -525,27 +370,37 @@ class DocumentQueryHelper:

370

371 def __init__(self, agent: Agent):

372 self.agent = agent

528 - self.store: DocumentQueryStore = asyncio.run(DocumentQueryStore.get(agent))

373 + self.store = DocumentQueryStore.get(agent)

374

530 - async def document_qa(self, document_uri: str, questions: Sequence[str]) -> Tuple[bool, str]:

531 - _ = await self.document_get_content(document_uri)

375 + async def document_qa(

376 + self, document_uri: str, questions: Sequence[str]

377 + ) -> Tuple[bool, str]:

378 + # index document

379 + _ = await self.document_get_content(document_uri, True)

380 content = ""

381 for question in questions:

382 human_content = f'Search Query: "{question}"'

535 - system_content = self.agent.parse_prompt("fw.document_query.optmimize_query.md")

536 -

537 - optimized_query = await self.agent.call_utility_model(

538 - system=system_content,

539 - message=human_content

383 + system_content = self.agent.parse_prompt(

384 + "fw.document_query.optmimize_query.md"

385 )

386

387 + optimized_query = (

388 + await self.agent.call_utility_model(

389 + system=system_content, message=human_content

390 + )

391 + ).strip()

392 +

393 + normalized_uri = self.store.normalize_uri(document_uri)

394 chunks = await self.store.search_document(

543 - document_uri=document_uri,

544 - query=str(optimized_query),

545 - limit=10000,

546 - threshold=0.66

395 + document_uri=normalized_uri,

396 + query=optimized_query,

397 + limit=100,

398 + threshold=DEFAULT_SEARCH_THRESHOLD,

399 + )

400 + content += (

401 + "\n\n----\n\n".join([chunk.page_content for chunk in chunks])

402 + + "\n\n----\n\n"

403 )

548 - content += "\n\n----\n\n".join([chunk.page_content for chunk in chunks]) + "\n\n----\n\n"

404

405 if not content:

406 content = f"!!! No content found for document: {document_uri} matching queries: {json.dumps(questions)}"

@@ -553,19 +408,23 @@ class DocumentQueryHelper:

408

409 questions_str = "\n".join([f" * {question}" for question in questions])

410

556 - qa_system_message = self.agent.parse_prompt("fw.document_query.system_prompt.md")

411 + qa_system_message = self.agent.parse_prompt(

412 + "fw.document_query.system_prompt.md"

413 + )

414 qa_user_message = f"# Document:\n{content}\n\n# Queries:\n{questions_str}"

415

416 ai_response = await self.agent.call_chat_model(

560 - prompt=ChatPromptTemplate.from_messages([

561 - SystemMessage(content=qa_system_message),

562 - HumanMessage(content=qa_user_message),

563 - ])

417 + prompt=ChatPromptTemplate.from_messages(

418 + [

419 + SystemMessage(content=qa_system_message),

420 + HumanMessage(content=qa_user_message),

421 + ]

422 + )

423 )

424

425 return True, str(ai_response)

426

568 - async def document_get_content(self, document_uri: str) -> str:

427 + async def document_get_content(self, document_uri: str, add_to_db: bool = False) -> str:

428 url = urlparse(document_uri)

429 scheme = url.scheme or "file"

430 mimetype, encoding = mimetypes.guess_type(document_uri)

@@ -579,7 +438,11 @@ class DocumentQueryHelper:

438 while not response and retries < 3:

439 try:

440 async with aiohttp.ClientSession() as session:

582 - response = await session.head(document_uri, timeout=aiohttp.ClientTimeout(total=2.0), allow_redirects=True)

441 + response = await session.head(

442 + document_uri,

443 + timeout=aiohttp.ClientTimeout(total=2.0),

444 + allow_redirects=True,

445 + )

446 if response.status > 399:

447 raise Exception(response.status)

448 break

@@ -589,32 +452,41 @@ class DocumentQueryHelper:

452 retries += 1

453

454 if not response:

592 - raise ValueError(f"DocumentQueryHelper::document_get_content: Document fetch error: {document_uri} ({last_error})")

455 + raise ValueError(

456 + f"DocumentQueryHelper::document_get_content: Document fetch error: {document_uri} ({last_error})"

457 + )

458

459 mimetype = response.headers["content-type"]

460 if "content-length" in response.headers:

596 - content_length = float(response.headers["content-length"]) / 1024 / 1024 # MB

461 + content_length = (

462 + float(response.headers["content-length"]) / 1024 / 1024

463 + ) # MB

464 if content_length > 50.0:

598 - raise ValueError(f"Document content length exceeds max. 50MB: {content_length} MB ({document_uri})")

599 - if mimetype and '; charset=' in mimetype:

600 - mimetype = mimetype.split('; charset=')[0]

465 + raise ValueError(

466 + f"Document content length exceeds max. 50MB: {content_length} MB ({document_uri})"

467 + )

468 + if mimetype and "; charset=" in mimetype:

469 + mimetype = mimetype.split("; charset=")[0]

470

471 if scheme == "file":

472 try:

604 - document_uri = os.path.abspath(url.path)

473 + document_uri = files.fix_dev_path(url.path)

474 except Exception as e:

475 raise ValueError(f"Invalid document path '{url.path}'") from e

476

477 if encoding:

609 - raise ValueError(f"Compressed documents are unsupported '{encoding}' ({document_uri})")

478 + raise ValueError(

479 + f"Compressed documents are unsupported '{encoding}' ({document_uri})"

480 + )

481

482 if mimetype == "application/octet-stream":

612 - raise ValueError(f"Unsupported document mimetype '{mimetype}' ({document_uri})")

483 + raise ValueError(

484 + f"Unsupported document mimetype '{mimetype}' ({document_uri})"

485 + )

486

487 # Use the store's normalization method

615 - document_uri_norm = self.store._normalize_uri(document_uri)

488 + document_uri_norm = self.store.normalize_uri(document_uri)

489

617 - await self.store.expire_documents(7)

490 exists = await self.store.document_exists(document_uri_norm)

491 document_content = ""

492 if not exists:

@@ -627,14 +499,19 @@ class DocumentQueryHelper:

499 elif mimetype == "application/pdf":

500 document_content = self.handle_pdf_document(document_uri, scheme)

501 else:

630 - document_content = self.handle_unstructured_document(document_uri, scheme)

631 - await self.store.add_document(document_content, document_uri_norm)

502 + document_content = self.handle_unstructured_document(

503 + document_uri, scheme

504 + )

505 + if add_to_db:

506 + await self.store.add_document(document_content, document_uri_norm)

507 else:

508 doc = await self.store.get_document(document_uri_norm)

509 if doc:

510 document_content = doc.page_content

511 else:

637 - raise ValueError(f"DocumentQueryHelper::document_get_content: Document not found: {document_uri_norm}")

512 + raise ValueError(

513 + f"DocumentQueryHelper::document_get_content: Document not found: {document_uri_norm}"

514 + )

515 return document_content

516

517 def handle_image_document(self, document: str, scheme: str) -> str:

@@ -646,14 +523,19 @@ class DocumentQueryHelper:

523 parts: list[Document] = loader.load()

524 elif scheme == "file":

525 # Use RFC file operations instead of TextLoader

649 - file_content_bytes = rfc_files.read_file_bin(document)

650 - file_content = file_content_bytes.decode('utf-8')

526 + file_content_bytes = files.read_file_bin(document)

527 + file_content = file_content_bytes.decode("utf-8")

528 # Create Document manually since we're not using TextLoader

529 parts = [Document(page_content=file_content, metadata={"source": document})]

530 else:

531 raise ValueError(f"Unsupported scheme: {scheme}")

532

656 - return "\n".join([element.page_content for element in MarkdownifyTransformer().transform_documents(parts)])

533 + return "\n".join(

534 + [

535 + element.page_content

536 + for element in MarkdownifyTransformer().transform_documents(parts)

537 + ]

538 + )

539

540 def handle_text_document(self, document: str, scheme: str) -> str:

541 if scheme in ["http", "https"]:

@@ -661,10 +543,12 @@ class DocumentQueryHelper:

543 elements: list[Document] = loader.load()

544 elif scheme == "file":

545 # Use RFC file operations instead of TextLoader

664 - file_content_bytes = rfc_files.read_file_bin(document)

665 - file_content = file_content_bytes.decode('utf-8')

546 + file_content_bytes = files.read_file_bin(document)

547 + file_content = file_content_bytes.decode("utf-8")

548 # Create Document manually since we're not using TextLoader

667 - elements = [Document(page_content=file_content, metadata={"source": document})]

549 + elements = [

550 + Document(page_content=file_content, metadata={"source": document})

551 + ]

552 else:

553 raise ValueError(f"Unsupported scheme: {scheme}")

554

@@ -674,27 +558,33 @@ class DocumentQueryHelper:

558 temp_file_path = ""

559 if scheme == "file":

560 # Use RFC file operations to read the PDF file as binary

677 - file_content_bytes = rfc_files.read_file_bin(document)

561 + file_content_bytes = files.read_file_bin(document)

562 # Create a temporary file for PyMuPDFLoader since it needs a file path

563 import tempfile

680 - with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:

564 +

565 + with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:

566 temp_file.write(file_content_bytes)

567 temp_file_path = temp_file.name

568 elif scheme in ["http", "https"]:

569 # download the file from the web url to a temporary file using python libraries for downloading

570 import requests

571 import tempfile

687 - with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:

572 +

573 + with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:

574 response = requests.get(document, timeout=10.0)

575 if response.status_code != 200:

690 - raise ValueError(f"DocumentQueryHelper::handle_pdf_document: Failed to download PDF from {document}: {response.status_code}")

576 + raise ValueError(

577 + f"DocumentQueryHelper::handle_pdf_document: Failed to download PDF from {document}: {response.status_code}"

578 + )

579 temp_file.write(response.content)

580 temp_file_path = temp_file.name

581 else:

582 raise ValueError(f"Unsupported scheme: {scheme}")

583

584 if not os.path.exists(temp_file_path):

697 - raise ValueError(f"DocumentQueryHelper::handle_pdf_document: Temporary file not found: {temp_file_path}")

585 + raise ValueError(

586 + f"DocumentQueryHelper::handle_pdf_document: Temporary file not found: {temp_file_path}"

587 + )

588

589 try:

590 try:

@@ -710,17 +600,21 @@ class DocumentQueryHelper:

600 elements: list[Document] = loader.load()

601 contents = "\n".join([element.page_content for element in elements])

602 except Exception as e:

713 - PrintStyle.error(f"DocumentQueryHelper::handle_pdf_document: Error loading with PyMuPDF: {e}")

603 + PrintStyle.error(

604 + f"DocumentQueryHelper::handle_pdf_document: Error loading with PyMuPDF: {e}"

605 + )

606 contents = ""

607

608 if not contents:

609 import pdf2image

610 import pytesseract

611

720 - PrintStyle.debug(f"DocumentQueryHelper::handle_pdf_document: FALLBACK Converting PDF to images: {temp_file_path}")

612 + PrintStyle.debug(

613 + f"DocumentQueryHelper::handle_pdf_document: FALLBACK Converting PDF to images: {temp_file_path}"

614 + )

615

616 # Convert PDF to images

723 - pages = pdf2image.convert_from_path(temp_file_path) # type: ignore

617 + pages = pdf2image.convert_from_path(temp_file_path) # type: ignore

618 for page in pages:

619 contents += pytesseract.image_to_string(page) + "\n\n"

620

@@ -742,10 +636,11 @@ class DocumentQueryHelper:

636 elements = loader.load()

637 elif scheme == "file":

638 # Use RFC file operations to read the file as binary

745 - file_content_bytes = rfc_files.read_file_bin(document)

639 + file_content_bytes = files.read_file_bin(document)

640 # Create a temporary file for UnstructuredLoader since it needs a file path

641 import tempfile

642 import os

643 +

644 # Get file extension to preserve it for proper processing

645 _, ext = os.path.splitext(document)

646 with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as temp_file:

python/helpers/files.py

@@ -253,8 +253,16 @@ def make_dirs(relative_path: str):

253

254

255 def get_abs_path(*relative_paths):

256 + "Convert relative paths to absolute paths based on the base directory."

257 return os.path.join(get_base_dir(), *relative_paths)

258

259 +def fix_dev_path(path:str):

260 + "On dev environment, convert /a0/... paths to local absolute paths"

261 + from python.helpers.runtime import is_development

262 + if is_development():

263 + if path.startswith("/a0/"):

264 + path = path.replace("/a0/", "")

265 + return get_abs_path(path)

266

267 def exists(*relative_paths):

268 path = get_abs_path(*relative_paths)

python/helpers/vector_db.py

+53 -17

@@ -27,27 +27,40 @@ class MyFaiss(FAISS):

27 async def aget_by_ids(self, ids: Sequence[str], /) -> List[Document]:

28 return self.get_by_ids(ids)

30 + def get_all_docs(self) -> dict[str, Document]:

31 + return self.docstore._dict # type: ignore

32 +

34 class VectorDB:

32 - def __init__(self, agent: Agent):

33 - self.agent = agent

34 - self.store = InMemoryByteStore()

35 - self.model = agent.get_embedding_model()

36 -

37 - self.embedder = CacheBackedEmbeddings.from_bytes_store(

38 - self.model,

39 - self.store,

40 - namespace=getattr(

41 - self.model,

42 - "model",

43 - getattr(self.model, "model_name", "default"),

44 - ),

35 +

36 + _cached_embeddings: dict[str, CacheBackedEmbeddings] = {}

37 +

38 + @staticmethod

39 + def _get_embeddings(agent: Agent):

40 + model = agent.get_embedding_model()

41 + namespace = getattr(

42 + model,

43 + "model",

44 + getattr(model, "model_name", "default"),

45 )

46 + if namespace not in VectorDB._cached_embeddings:

47 + store = InMemoryByteStore()

48 + VectorDB._cached_embeddings[namespace] = (

49 + CacheBackedEmbeddings.from_bytes_store(

50 + model,

51 + store,

52 + namespace=namespace,

53 + )

54 + )

55 + return VectorDB._cached_embeddings[namespace]

47 - self.index = faiss.IndexFlatIP(len(self.embedder.embed_query("example")))

57 + def __init__(self, agent: Agent):

58 + self.agent = agent

59 + self.embeddings = self._get_embeddings(agent)

60 + self.index = faiss.IndexFlatIP(len(self.embeddings.embed_query("example")))

62 self.db = MyFaiss(

50 - embedding_function=self.embedder,

63 + embedding_function=self.embeddings,

64 index=self.index,

65 docstore=InMemoryDocstore(),

66 index_to_docstore_id={},

@@ -56,7 +69,7 @@ class VectorDB:

69 relevance_score_fn=cosine_normalizer,

70 )

59 - async def search_similarity_threshold(

72 + async def search_by_similarity_threshold(

73 self, query: str, limit: int, threshold: float, filter: str = ""

74 ):

75 comparator = get_comparator(filter) if filter else None

@@ -74,6 +87,18 @@ class VectorDB:

87 filter=comparator,

88 )

90 + async def search_by_metadata(self, filter: str, limit: int = 0) -> list[Document]:

91 + comparator = get_comparator(filter)

92 + all_docs = self.db.get_all_docs()

93 + result = []

94 + for doc in all_docs.values():

95 + if comparator(doc.metadata):

96 + result.append(doc)

97 + # stop if limit reached and limit > 0

98 + if limit > 0 and len(result) >= limit:

99 + break

100 + return result

101 +

102 async def insert_documents(self, docs: list[Document]):

103 ids = [str(uuid.uuid4()) for _ in range(len(docs))]

104

@@ -90,6 +115,16 @@ class VectorDB:

115 self.db.add_documents(documents=docs, ids=ids)

116 return ids

117

118 + async def delete_documents_by_ids(self, ids: list[str]):

119 + # aget_by_ids is not yet implemented in faiss, need to do a workaround

120 + rem_docs = await self.db.aget_by_ids(

121 + ids

122 + ) # existing docs to remove (prevents error)

123 + if rem_docs:

124 + rem_ids = [doc.metadata["id"] for doc in rem_docs] # ids to remove

125 + await self.db.adelete(ids=rem_ids)

126 + return rem_docs

127 +

128

129 def format_docs_plain(docs: list[Document]) -> list[str]:

130 result = []

@@ -113,7 +148,8 @@ def cosine_normalizer(val: float) -> float:

148 def get_comparator(condition: str):

149 def comparator(data: dict[str, Any]):

150 try:

116 - return eval(condition, {}, data)

151 + result = eval(condition, {}, data)

152 + return result

153 except Exception as e:

154 # PrintStyle.error(f"Error evaluating condition: {e}")

155 return False