Spaces:
Running
Running
File size: 3,319 Bytes
9fa4cb8 13a9ae5 fefceb2 9fa4cb8 b568914 9fa4cb8 c58c522 fefceb2 9fa4cb8 b568914 9fa4cb8 c58c522 9fa4cb8 c58c522 9fa4cb8 fefceb2 13a9ae5 9fa4cb8 7e612d6 9fa4cb8 e5a699f 9fa4cb8 13a9ae5 9fa4cb8 c58c522 9fa4cb8 c58c522 7e612d6 9fa4cb8 fefceb2 9fa4cb8 c58c522 b568914 9fa4cb8 b568914 9fa4cb8 b568914 9fa4cb8 b568914 9fa4cb8 b568914 9fa4cb8 b568914 9fa4cb8 c58c522 9fa4cb8 | 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 | import os
from typing import Type, Optional, Any
from pydantic import BaseModel, Field
from markitdown import MarkItDown
from chonkie import SemanticChunker
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from fastembed import TextEmbedding
try:
from crewai.tools import BaseTool
except ImportError:
from langchain.tools import BaseTool
_EMBEDDER = None
def _get_embedder():
global _EMBEDDER
if _EMBEDDER is None:
_EMBEDDER = TextEmbedding("BAAI/bge-small-en-v1.5")
return _EMBEDDER
class SearchInput(BaseModel):
query: str = Field(..., description="Search query")
class DocumentSearchTool(BaseTool):
name: str = "DocumentSearchTool"
description: str = "Search uploaded document for relevant passages."
args_schema: Type[BaseModel] = SearchInput
file_path: Optional[str] = None
client: Optional[Any] = None
COLLECTION : str = "neuraldocs_collection"
EMBED_MODEL : str = "minishlab/potion-base-8M"
CHUNK_SIZE : int = 128
SIMILARITY_T: float = 0.5
TOP_K : int = 2
SEPARATOR : str = "\n---\n"
def __init__(self, file_path: str):
super().__init__(file_path=file_path, client=QdrantClient(":memory:"))
self._build_index()
def _to_text(self) -> str:
converter = MarkItDown()
result = converter.convert(self.file_path)
text = result.text_content.strip()
if not text:
raise ValueError(f"Could not extract text from '{self.file_path}'.")
return text[:5000]
def _chunk(self, text: str) -> list:
chunker = SemanticChunker(
embedding_model=self.EMBED_MODEL,
threshold=self.SIMILARITY_T,
chunk_size=self.CHUNK_SIZE,
min_sentences=1,
)
return [c.text for c in chunker.chunk(text) if c.text.strip()]
def _build_index(self) -> None:
chunks = self._chunk(self._to_text())
if not chunks:
raise ValueError("No text chunks were produced from the document.")
source_name = os.path.basename(self.file_path)
vectors = [v.tolist() for v in _get_embedder().embed(chunks)]
self.client.create_collection(
collection_name=self.COLLECTION,
vectors_config=VectorParams(size=len(vectors[0]), distance=Distance.COSINE),
)
self.client.upsert(
collection_name=self.COLLECTION,
points=[
PointStruct(
id=i,
vector=vectors[i],
payload={"document": chunks[i], "source": source_name, "chunk_id": i},
)
for i in range(len(chunks))
],
)
def _run(self, query: str) -> str:
qvec = next(iter(_get_embedder().embed([query]))).tolist()
res = self.client.query_points(
collection_name=self.COLLECTION,
query=qvec,
limit=self.TOP_K,
)
passages = [
p.payload.get("document", "") for p in res.points if p.payload
]
passages = [t for t in passages if t and t.strip()]
if not passages:
return "No relevant passages found in the document."
return self.SEPARATOR.join(passages)
|