chore: 添加虚拟环境到仓库
- 添加 backend_service/venv 虚拟环境 - 包含所有Python依赖包 - 注意:虚拟环境约393MB,包含12655个文件
This commit is contained in:
@@ -0,0 +1,37 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The retrieval-augmented generation (RAG) module in AgentScope."""
|
||||
|
||||
from ._document import (
|
||||
DocMetadata,
|
||||
Document,
|
||||
)
|
||||
from ._reader import (
|
||||
ReaderBase,
|
||||
TextReader,
|
||||
PDFReader,
|
||||
ImageReader,
|
||||
WordReader,
|
||||
)
|
||||
from ._store import (
|
||||
VDBStoreBase,
|
||||
QdrantStore,
|
||||
MilvusLiteStore,
|
||||
)
|
||||
from ._knowledge_base import KnowledgeBase
|
||||
from ._simple_knowledge import SimpleKnowledge
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ReaderBase",
|
||||
"TextReader",
|
||||
"PDFReader",
|
||||
"ImageReader",
|
||||
"WordReader",
|
||||
"DocMetadata",
|
||||
"Document",
|
||||
"VDBStoreBase",
|
||||
"QdrantStore",
|
||||
"MilvusLiteStore",
|
||||
"KnowledgeBase",
|
||||
"SimpleKnowledge",
|
||||
]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,51 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The document data structure used in RAG as the data chunk and
|
||||
retrieval result."""
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import shortuuid
|
||||
from dashscope.api_entities.dashscope_response import DictMixin
|
||||
|
||||
from ..message import (
|
||||
TextBlock,
|
||||
ImageBlock,
|
||||
VideoBlock,
|
||||
)
|
||||
from ..types import Embedding
|
||||
|
||||
|
||||
@dataclass
|
||||
class DocMetadata(DictMixin):
|
||||
"""The metadata of the document."""
|
||||
|
||||
content: TextBlock | ImageBlock | VideoBlock
|
||||
"""The data content, e.g., text, image, video."""
|
||||
|
||||
doc_id: str
|
||||
"""The document ID."""
|
||||
|
||||
chunk_id: int
|
||||
"""The chunk ID."""
|
||||
|
||||
total_chunks: int
|
||||
"""The total number of chunks."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""The data chunk."""
|
||||
|
||||
metadata: DocMetadata
|
||||
"""The metadata of the data chunk."""
|
||||
|
||||
id: str = field(default_factory=shortuuid.uuid)
|
||||
"""The unique ID of the data chunk."""
|
||||
|
||||
# The fields that will be filled when the document is added to or
|
||||
# retrieved from the knowledge base.
|
||||
|
||||
embedding: Embedding | None = field(default_factory=lambda: None)
|
||||
"""The embedding of the data chunk."""
|
||||
|
||||
score: float | None = None
|
||||
"""The relevance score of the data chunk."""
|
||||
@@ -0,0 +1,130 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The knowledge base abstraction for retrieval-augmented generation (RAG)."""
|
||||
from abc import abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from ._reader import Document
|
||||
from ..embedding import EmbeddingModelBase
|
||||
from ._store import VDBStoreBase
|
||||
from ..message import TextBlock
|
||||
from ..tool import ToolResponse
|
||||
|
||||
|
||||
class KnowledgeBase:
|
||||
"""The knowledge base abstraction for retrieval-augmented generation
|
||||
(RAG).
|
||||
|
||||
The ``retrieve`` and ``add_documents`` methods need to be implemented
|
||||
in the subclasses. We also provide a quick method ``retrieve_knowledge``
|
||||
that enables the agent to retrieve knowledge easily.
|
||||
"""
|
||||
|
||||
embedding_store: VDBStoreBase
|
||||
"""The embedding store for the knowledge base."""
|
||||
|
||||
embedding_model: EmbeddingModelBase
|
||||
"""The embedding model for the knowledge base."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedding_store: VDBStoreBase,
|
||||
embedding_model: EmbeddingModelBase,
|
||||
) -> None:
|
||||
"""Initialize the knowledge base."""
|
||||
self.embedding_store = embedding_store
|
||||
self.embedding_model = embedding_model
|
||||
|
||||
@abstractmethod
|
||||
async def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
limit: int = 5,
|
||||
score_threshold: float | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Document]:
|
||||
"""Retrieve relevant documents by the given query.
|
||||
|
||||
Args:
|
||||
query (`str`):
|
||||
The query string to retrieve relevant documents.
|
||||
limit (`int`, defaults to 5):
|
||||
The number of relevant documents to retrieve.
|
||||
score_threshold (`float | None`, defaults to `None`):
|
||||
The score threshold to filter the retrieved documents. If
|
||||
provided, only documents with a score higher than the
|
||||
threshold will be returned.
|
||||
**kwargs (`Any`):
|
||||
Other keyword arguments for the vector database search API.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_documents(
|
||||
self,
|
||||
documents: list[Document],
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Add documents to the knowledge base, which will embed the documents
|
||||
and store them in the embedding store.
|
||||
|
||||
Args:
|
||||
documents (`list[Document]`):
|
||||
A list of documents to add.
|
||||
"""
|
||||
|
||||
# A quick method that enable the agent to retrieve knowledge
|
||||
# Developers can wrap the `retrieve` method by themselves to support
|
||||
# more flexible usage
|
||||
async def retrieve_knowledge(
|
||||
self,
|
||||
query: str,
|
||||
limit: int = 5,
|
||||
score_threshold: float | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ToolResponse:
|
||||
"""Retrieve relevant documents from the knowledge base. Note the
|
||||
`query` parameter is directly related to the retrieval quality, and
|
||||
for the same question, you can try many different queries to get the
|
||||
best results. Adjust the `limit` and `score_threshold` parameters
|
||||
to get more or fewer results.
|
||||
|
||||
Args:
|
||||
query (`str`):
|
||||
The query string, which should be specific and concise. For
|
||||
example, you should provide the specific name instead of
|
||||
"you", "my", "he", "she", etc.
|
||||
limit (`int`, defaults to 3):
|
||||
The number of relevant documents to retrieve.
|
||||
score_threshold (`float`, defaults to 0.8):
|
||||
A threshold in [0, 1] and only the relevance score above this
|
||||
threshold will be returned. Reduce this value to get more
|
||||
results.
|
||||
"""
|
||||
|
||||
docs = await self.retrieve(
|
||||
query=query,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if len(docs):
|
||||
return ToolResponse(
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=f"Score: {_.score}, "
|
||||
f"Content: {_.metadata.content['text']}",
|
||||
)
|
||||
for _ in docs
|
||||
],
|
||||
)
|
||||
return ToolResponse(
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text="No relevant documents found. TRY to reduce the "
|
||||
"`score_threshold` parameter to get "
|
||||
"more results.",
|
||||
),
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,18 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The reader abstraction for retrieval-augmented generation (RAG)."""
|
||||
|
||||
from ._reader_base import ReaderBase, Document
|
||||
from ._text_reader import TextReader
|
||||
from ._pdf_reader import PDFReader
|
||||
from ._word_reader import WordReader
|
||||
from ._image_reader import ImageReader
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Document",
|
||||
"ReaderBase",
|
||||
"TextReader",
|
||||
"PDFReader",
|
||||
"WordReader",
|
||||
"ImageReader",
|
||||
]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,67 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The Image reader modules"""
|
||||
import hashlib
|
||||
|
||||
from .. import DocMetadata
|
||||
from ...message import ImageBlock, URLSource
|
||||
from .._reader import ReaderBase, Document
|
||||
|
||||
|
||||
class ImageReader(ReaderBase):
|
||||
"""A simple image reader that wraps the image into a Document object.
|
||||
|
||||
This class is only a simple implementation to support multimodal RAG.
|
||||
"""
|
||||
|
||||
async def __call__(self, image_url: str | list[str]) -> list[Document]:
|
||||
"""Read an image and return the wrapped Document object.
|
||||
|
||||
Args:
|
||||
image_url (`str | list[str]`):
|
||||
The image URL(s) or path(s).
|
||||
|
||||
Returns:
|
||||
`list[Document]`:
|
||||
A list of Document objects containing the image data.
|
||||
"""
|
||||
# Read the image data and wrap it into a Document object.
|
||||
if isinstance(image_url, str):
|
||||
image_url = [image_url]
|
||||
|
||||
image_blocks: list[ImageBlock] = [
|
||||
ImageBlock(
|
||||
type="image",
|
||||
source=URLSource(
|
||||
type="url",
|
||||
url=_,
|
||||
),
|
||||
)
|
||||
for _ in image_url
|
||||
]
|
||||
|
||||
doc_idx = [self.get_doc_id(_) for _ in image_url]
|
||||
|
||||
return [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=image_block,
|
||||
doc_id=doc_id,
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
)
|
||||
for doc_id, image_block in zip(doc_idx, image_blocks)
|
||||
]
|
||||
|
||||
def get_doc_id(self, image_path: str) -> str:
|
||||
"""Generate a document ID based on the image path.
|
||||
|
||||
Args:
|
||||
image_path (`str`):
|
||||
The image path or URL.
|
||||
|
||||
Returns:
|
||||
`str`:
|
||||
The generated document ID.
|
||||
"""
|
||||
return hashlib.md5(image_path.encode("utf-8")).hexdigest()
|
||||
@@ -0,0 +1,86 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The PDF reader to read and chunk PDF files."""
|
||||
import hashlib
|
||||
from typing import Literal
|
||||
|
||||
from ._reader_base import ReaderBase
|
||||
from ._text_reader import TextReader
|
||||
from .._document import Document
|
||||
|
||||
|
||||
class PDFReader(ReaderBase):
|
||||
"""The PDF reader that splits text into chunks by a fixed chunk size."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: int = 512,
|
||||
split_by: Literal["char", "sentence", "paragraph"] = "sentence",
|
||||
) -> None:
|
||||
"""Initialize the text reader.
|
||||
|
||||
Args:
|
||||
chunk_size (`int`, default to 512):
|
||||
The size of each chunk, in number of characters.
|
||||
split_by (`Literal["char", "sentence", "paragraph"]`, default to \
|
||||
"sentence"):
|
||||
The unit to split the text, can be "char", "sentence", or
|
||||
"paragraph". The "sentence" option is implemented using the
|
||||
"nltk" library, which only supports English text.
|
||||
"""
|
||||
if chunk_size <= 0:
|
||||
raise ValueError(
|
||||
f"The chunk_size must be positive, got {chunk_size}",
|
||||
)
|
||||
|
||||
if split_by not in ["char", "sentence", "paragraph"]:
|
||||
raise ValueError(
|
||||
"The split_by must be one of 'char', 'sentence' or "
|
||||
f"'paragraph', got {split_by}",
|
||||
)
|
||||
|
||||
self.chunk_size = chunk_size
|
||||
self.split_by = split_by
|
||||
|
||||
# To avoid code duplication, we use TextReader to do the chunking.
|
||||
self._text_reader = TextReader(
|
||||
self.chunk_size,
|
||||
self.split_by,
|
||||
)
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
pdf_path: str,
|
||||
) -> list[Document]:
|
||||
"""Read a PDF file, split it into chunks, and return a list of
|
||||
Document objects.
|
||||
|
||||
Args:
|
||||
pdf_path (`str`):
|
||||
The input PDF file path.
|
||||
"""
|
||||
try:
|
||||
from pypdf import PdfReader
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Please install pypdf to use the PDF reader. "
|
||||
"You can install it by `pip install pypdf`.",
|
||||
) from e
|
||||
|
||||
reader = PdfReader(pdf_path)
|
||||
|
||||
gather_texts = []
|
||||
for page in reader.pages:
|
||||
gather_texts.append(page.extract_text())
|
||||
|
||||
doc_id = hashlib.sha256(pdf_path.encode("utf-8")).hexdigest()
|
||||
|
||||
docs = await self._text_reader("\n\n".join(gather_texts))
|
||||
for doc in docs:
|
||||
doc.id = doc_id
|
||||
|
||||
return docs
|
||||
|
||||
def get_doc_id(self, pdf_path: str) -> str:
|
||||
"""Get the document ID. This function can be used to check if the
|
||||
doc_id already exists in the knowledge base."""
|
||||
return hashlib.sha256(pdf_path.encode("utf-8")).hexdigest()
|
||||
@@ -0,0 +1,27 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The reader base class for retrieval-augmented generation (RAG)."""
|
||||
from abc import abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from .._document import Document
|
||||
|
||||
|
||||
class ReaderBase:
|
||||
"""The reader base class, which is responsible for reading the original
|
||||
data, splitting it into chunks, and converting each chunk into a `Document`
|
||||
object."""
|
||||
|
||||
@abstractmethod
|
||||
async def __call__(self, *args: Any, **kwargs: Any) -> list[Document]:
|
||||
"""The async call function that takes the input files and returns the
|
||||
vector records"""
|
||||
|
||||
@abstractmethod
|
||||
def get_doc_id(self, *args: Any, **kwargs: Any) -> str:
|
||||
"""Get a unique document ID for the input data. This method is to
|
||||
expose the document ID generation logic to the developers
|
||||
|
||||
Returns:
|
||||
`str`:
|
||||
A unique document ID for the input data.
|
||||
"""
|
||||
@@ -0,0 +1,149 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The text reader that reads text into vector records."""
|
||||
import hashlib
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
from ._reader_base import ReaderBase, Document
|
||||
from .._document import DocMetadata
|
||||
from ..._logging import logger
|
||||
from ...message import TextBlock
|
||||
|
||||
|
||||
class TextReader(ReaderBase):
|
||||
"""The text reader that splits text into chunks by a fixed chunk size
|
||||
and chunk overlap."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: int = 512,
|
||||
split_by: Literal["char", "sentence", "paragraph"] = "sentence",
|
||||
) -> None:
|
||||
"""Initialize the text reader.
|
||||
|
||||
Args:
|
||||
chunk_size (`int`, default to 512):
|
||||
The size of each chunk, in number of characters.
|
||||
split_by (`Literal["char", "paragraph"]`, default to \
|
||||
"sentence"):
|
||||
The unit to split the text, can be "char", "sentence", or
|
||||
"paragraph". Note that "sentence" is implemented by "nltk"
|
||||
library, which only supports English text.
|
||||
"""
|
||||
if chunk_size <= 0:
|
||||
raise ValueError(
|
||||
f"The chunk_size must be positive, got {chunk_size}",
|
||||
)
|
||||
|
||||
if split_by not in ["char", "sentence", "paragraph"]:
|
||||
raise ValueError(
|
||||
"The split_by must be one of 'char', 'sentence' or "
|
||||
f"'paragraph', got {split_by}",
|
||||
)
|
||||
|
||||
self.chunk_size = chunk_size
|
||||
self.split_by = split_by
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
text: str,
|
||||
) -> list[Document]:
|
||||
"""Read a text string, split it into chunks, and return a list of
|
||||
Document objects.
|
||||
|
||||
Args:
|
||||
text (`str`):
|
||||
The input text string, or a path to the local text file.
|
||||
|
||||
Returns:
|
||||
`list[Document]`:
|
||||
A list of Document objects, where the metadata contains the
|
||||
chunked text, doc id and chunk id.
|
||||
"""
|
||||
if os.path.exists(text) and os.path.isfile(text):
|
||||
logger.info("Reading text from local file: %s", text)
|
||||
with open(text, "r", encoding="utf-8") as file:
|
||||
text = file.read()
|
||||
|
||||
logger.info(
|
||||
"Reading text with chunk_size=%d, split_by=%s",
|
||||
self.chunk_size,
|
||||
self.split_by,
|
||||
)
|
||||
splits = []
|
||||
if self.split_by == "char":
|
||||
# Split by character
|
||||
for i in range(0, len(text), self.chunk_size):
|
||||
start = max(0, i)
|
||||
end = min(i + self.chunk_size, len(text))
|
||||
splits.append(text[start:end])
|
||||
|
||||
elif self.split_by == "sentence":
|
||||
try:
|
||||
import nltk
|
||||
|
||||
nltk.download("punkt", quiet=True)
|
||||
nltk.download("punkt_tab", quiet=True)
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"nltk is not installed. Please install it with "
|
||||
"`pip install nltk`.",
|
||||
) from e
|
||||
|
||||
sentences = nltk.sent_tokenize(text)
|
||||
|
||||
# Handle the chunk_size for sentences
|
||||
processed_sentences = []
|
||||
for _ in sentences:
|
||||
if len(_) <= self.chunk_size:
|
||||
processed_sentences.append(_)
|
||||
else:
|
||||
# If the sentence itself exceeds chunk size, we need to
|
||||
# truncate it
|
||||
chunks = [
|
||||
_[j : j + self.chunk_size]
|
||||
for j in range(0, len(_), self.chunk_size)
|
||||
]
|
||||
processed_sentences.extend(chunks)
|
||||
|
||||
splits.extend(processed_sentences)
|
||||
|
||||
elif self.split_by == "paragraph":
|
||||
paragraphs = [_ for _ in text.split("\n") if len(_)]
|
||||
for para in paragraphs:
|
||||
if len(para) <= self.chunk_size:
|
||||
splits.append(para)
|
||||
|
||||
else:
|
||||
# If the paragraph itself exceeds chunk size, we need to
|
||||
# truncate it
|
||||
chunks = [
|
||||
para[k : k + self.chunk_size]
|
||||
for k in range(0, len(para), self.chunk_size)
|
||||
]
|
||||
splits.extend(chunks)
|
||||
|
||||
logger.info(
|
||||
"Finished splitting the text into %d chunks.",
|
||||
len(splits),
|
||||
)
|
||||
|
||||
doc_id = self.get_doc_id(text)
|
||||
|
||||
return [
|
||||
Document(
|
||||
id=doc_id,
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(type="text", text=_),
|
||||
doc_id=doc_id,
|
||||
chunk_id=idx,
|
||||
total_chunks=len(splits),
|
||||
),
|
||||
)
|
||||
for idx, _ in enumerate(splits)
|
||||
]
|
||||
|
||||
def get_doc_id(self, text: str) -> str:
|
||||
"""Get the document ID. This function can be used to check if the
|
||||
doc_id already exists in the knowledge base."""
|
||||
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||||
@@ -0,0 +1,508 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# pylint: disable=W0212
|
||||
"""The Word reader to read and chunk Word documents."""
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
from typing import Literal, TYPE_CHECKING
|
||||
|
||||
|
||||
from ._reader_base import ReaderBase
|
||||
from ._text_reader import TextReader
|
||||
from .._document import Document, DocMetadata
|
||||
from ..._logging import logger
|
||||
from ...message import ImageBlock, Base64Source, TextBlock
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from docx.table import Table as DocxTable
|
||||
from docx.text.paragraph import Paragraph as DocxParagraph
|
||||
else:
|
||||
DocxTable = "docx.table.Table"
|
||||
DocxParagraph = "docx.text.paragraph.Paragraph"
|
||||
|
||||
|
||||
def _extract_text_from_paragraph(para: DocxParagraph) -> str:
|
||||
"""Extract text from a paragraph, including text in text boxes and shapes.
|
||||
|
||||
Args:
|
||||
para (`Paragraph`):
|
||||
The paragraph object from which to extract text.
|
||||
|
||||
|
||||
Returns:
|
||||
`str`:
|
||||
Extracted text
|
||||
"""
|
||||
text = ""
|
||||
|
||||
# Method 1: Extract all w:t elements directly from XML
|
||||
# (handles revisions, hyperlinks, etc.)
|
||||
from docx.oxml.ns import qn
|
||||
|
||||
for t_elem in para._element.findall(".//" + qn("w:t")):
|
||||
if t_elem.text:
|
||||
text += t_elem.text
|
||||
|
||||
# Method 2: If no text found, try standard text property
|
||||
if not text:
|
||||
text = para.text.strip()
|
||||
|
||||
# Method 3: If still no text, try to extract from text boxes and shapes
|
||||
if not text:
|
||||
# Check for text boxes (txbxContent)
|
||||
txbx_contents = para._element.findall(".//" + qn("w:txbxContent"))
|
||||
for txbx in txbx_contents:
|
||||
# Extract all text from paragraphs within the text box
|
||||
for p_elem in txbx.findall(".//" + qn("w:p")):
|
||||
for t_elem in p_elem.findall(".//" + qn("w:t")):
|
||||
if t_elem.text:
|
||||
text += t_elem.text
|
||||
|
||||
# Check for VML text boxes - use full namespace URI
|
||||
vml_ns = "{urn:schemas-microsoft-com:vml}"
|
||||
vml_textboxes = para._element.findall(".//" + vml_ns + "textbox")
|
||||
for vml_tb in vml_textboxes:
|
||||
for p_elem in vml_tb.findall(".//" + qn("w:p")):
|
||||
for t_elem in p_elem.findall(".//" + qn("w:t")):
|
||||
if t_elem.text:
|
||||
text += t_elem.text
|
||||
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _extract_table_data(table: DocxTable) -> list[list[str]]:
|
||||
"""Extract table data, handling merged cells and preserving line breaks
|
||||
within cells.
|
||||
|
||||
Args:
|
||||
table (`Table`):
|
||||
The table object from which to extract data.
|
||||
|
||||
Returns:
|
||||
`list[list[str]]`:
|
||||
Table data represented as a 2D list.
|
||||
"""
|
||||
|
||||
from docx.oxml.ns import qn
|
||||
|
||||
table_data = []
|
||||
# Extract table cell elements directly from XML
|
||||
for tr in table._element.findall(qn("w:tr")):
|
||||
row_data = []
|
||||
|
||||
tcs = tr.findall(qn("w:tc"))
|
||||
for tc in tcs:
|
||||
# Extract paragraphs within the table cell (preserve line breaks)
|
||||
paragraphs = []
|
||||
for p_elem in tc.findall(qn("w:p")):
|
||||
# Obtain all text elements within the paragraph
|
||||
texts = []
|
||||
for t_elem in p_elem.findall(".//" + qn("w:t")):
|
||||
if t_elem.text:
|
||||
texts.append(t_elem.text)
|
||||
|
||||
para_text = "".join(texts)
|
||||
if para_text:
|
||||
# Only add non-empty paragraphs
|
||||
paragraphs.append(para_text)
|
||||
|
||||
# Use \n to join multiple paragraphs
|
||||
cell_text = "\n".join(paragraphs)
|
||||
row_data.append(cell_text)
|
||||
|
||||
table_data.append(row_data)
|
||||
|
||||
return table_data
|
||||
|
||||
|
||||
def _extract_image_data(para: DocxParagraph) -> list[ImageBlock]:
|
||||
"""Extract image data from a paragraph.
|
||||
|
||||
Args:
|
||||
para (`Paragraph`):
|
||||
The paragraph object from which to extract images.
|
||||
|
||||
Returns:
|
||||
`list[ImageBlock]`:
|
||||
A list of image blocks with base64-encoded image data
|
||||
"""
|
||||
images = []
|
||||
|
||||
from docx.oxml.ns import qn
|
||||
|
||||
# Method 1: Find all drawing elements (modern Word format)
|
||||
drawings = para._element.findall(".//" + qn("w:drawing"))
|
||||
|
||||
for drawing in drawings:
|
||||
# Try to find blip elements (embedded images)
|
||||
blips = drawing.findall(".//" + qn("a:blip"))
|
||||
|
||||
for blip in blips:
|
||||
# Get the relationship ID
|
||||
embed = blip.get(qn("r:embed"))
|
||||
|
||||
if embed:
|
||||
try:
|
||||
# Get the image part from the document
|
||||
image_part = para.part.related_parts[embed]
|
||||
# Get the image binary data
|
||||
image_data = image_part.blob
|
||||
# Encode to base64
|
||||
image_base64 = base64.b64encode(image_data).decode("utf-8")
|
||||
|
||||
# Get image format from content type
|
||||
content_type = image_part.content_type
|
||||
|
||||
images.append(
|
||||
ImageBlock(
|
||||
type="image",
|
||||
source=Base64Source(
|
||||
type="base64",
|
||||
data=image_base64,
|
||||
media_type=content_type,
|
||||
),
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to extract image: %s",
|
||||
e,
|
||||
)
|
||||
|
||||
# Method 2: Check for pict elements (older Word format)
|
||||
picts = para._element.findall(".//" + qn("w:pict"))
|
||||
|
||||
for pict in picts:
|
||||
imagedatas = pict.findall(".//" + qn("v:imagedata"))
|
||||
|
||||
for imagedata in imagedatas:
|
||||
rel_id = imagedata.get(qn("r:id"))
|
||||
|
||||
if rel_id:
|
||||
try:
|
||||
image_part = para.part.related_parts[rel_id]
|
||||
image_data = image_part.blob
|
||||
image_base64 = base64.b64encode(image_data).decode("utf-8")
|
||||
|
||||
images.append(
|
||||
ImageBlock(
|
||||
type="image",
|
||||
source=Base64Source(
|
||||
type="base64",
|
||||
data=image_base64,
|
||||
media_type=image_part.content_type,
|
||||
),
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to extract image from pict: %s",
|
||||
e,
|
||||
)
|
||||
return images
|
||||
|
||||
|
||||
class WordReader(ReaderBase):
|
||||
"""The reader that supports reading text, image, and table content from
|
||||
Word documents (.docx files), and chunking the text content into smaller
|
||||
pieces.
|
||||
|
||||
.. note:: The table content is extracted in Markdown format.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: int = 512,
|
||||
split_by: Literal["char", "sentence", "paragraph"] = "sentence",
|
||||
include_image: bool = True,
|
||||
separate_table: bool = False,
|
||||
table_format: Literal["markdown", "json"] = "markdown",
|
||||
) -> None:
|
||||
"""Initialize the Word reader.
|
||||
|
||||
Args:
|
||||
chunk_size (`int`, default to 512):
|
||||
The size of each chunk, in number of characters.
|
||||
split_by (`Literal["char", "sentence", "paragraph"]`, default to \
|
||||
"sentence"):
|
||||
The unit to split the text, can be "char", "sentence", or
|
||||
"paragraph". The "sentence" option is implemented using the
|
||||
"nltk" library, which only supports English text.
|
||||
include_image (`bool`, default to False):
|
||||
Whether to include image content in the returned document. If
|
||||
activated, the embedding model you use must support image
|
||||
input, e.g. `DashScopeMultiModalEmbedding`.
|
||||
separate_table (`bool`, default to False):
|
||||
If True, tables will be treated as a new chunk to avoid
|
||||
truncation. But note when the table exceeds the chunk size,
|
||||
it will still be truncated.
|
||||
table_format (`Literal["markdown", "json"]`, \
|
||||
default to "markdown"):
|
||||
The format to extract table content. Note if the table cell
|
||||
contains `\n`, the Markdown format may not render correctly.
|
||||
In that case, you can use the `json` format, which extracts
|
||||
the table as a JSON string of a `list[list[str]]` object.
|
||||
"""
|
||||
if chunk_size <= 0:
|
||||
raise ValueError(
|
||||
f"The chunk_size must be positive, got {chunk_size}",
|
||||
)
|
||||
|
||||
if split_by not in ["char", "sentence", "paragraph"]:
|
||||
raise ValueError(
|
||||
"The split_by must be one of 'char', 'sentence' or "
|
||||
f"'paragraph', got {split_by}",
|
||||
)
|
||||
|
||||
if table_format not in ["markdown", "json"]:
|
||||
raise ValueError(
|
||||
"The table_format must be one of 'markdown' or 'json', "
|
||||
f"got {table_format}",
|
||||
)
|
||||
|
||||
self.chunk_size = chunk_size
|
||||
self.split_by = split_by
|
||||
self.include_image = include_image
|
||||
self.separate_table = separate_table
|
||||
self.table_format = table_format
|
||||
|
||||
# To avoid code duplication, we use TextReader to do the chunking.
|
||||
self._text_reader = TextReader(
|
||||
self.chunk_size,
|
||||
self.split_by,
|
||||
)
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
word_path: str,
|
||||
) -> list[Document]:
|
||||
"""Read a Word document, split it into chunks, and return a list of
|
||||
Document objects. The text, image, and table content will be returned
|
||||
in the same order as they appear in the Word document.
|
||||
|
||||
Args:
|
||||
word_path (`str`):
|
||||
The input Word document file path (.docx file).
|
||||
|
||||
Returns:
|
||||
`list[Document]`:
|
||||
A list of Document objects, where the metadata contains the
|
||||
chunked text, doc id and chunk id.
|
||||
"""
|
||||
|
||||
blocks = self._get_data_blocks(word_path)
|
||||
|
||||
doc_id = self.get_doc_id(word_path)
|
||||
documents = []
|
||||
for block in blocks:
|
||||
if block["type"] == "text":
|
||||
for _ in await self._text_reader(block["text"]):
|
||||
documents.append(
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=_.metadata.content,
|
||||
doc_id=doc_id,
|
||||
# The chunk_id and total_chunks will be reset
|
||||
chunk_id=0,
|
||||
total_chunks=0,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
elif block["type"] == "image":
|
||||
documents.append(
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=block,
|
||||
doc_id=doc_id,
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
# Set chunk ids and total chunks
|
||||
total_chunks = len(documents)
|
||||
for idx, doc in enumerate(documents):
|
||||
doc.metadata.chunk_id = idx
|
||||
doc.metadata.total_chunks = total_chunks
|
||||
|
||||
return documents
|
||||
|
||||
def _get_data_blocks(self, word_path: str) -> list[TextBlock | ImageBlock]:
|
||||
"""This function will return a list of dicts, each dict has a
|
||||
'type' field indicating 'text', 'table', or 'image', and a
|
||||
corresponding field containing the actual data.
|
||||
|
||||
Args:
|
||||
word_path (`str`):
|
||||
The input Word document file path (.docx file).
|
||||
|
||||
Returns:
|
||||
`list[TextBlock | ImageBlock]`:
|
||||
A list of data blocks extracted from the Word document.
|
||||
"""
|
||||
# Read the Word document
|
||||
try:
|
||||
from docx import Document as DocxDocument
|
||||
from docx.oxml import CT_P, CT_Tbl
|
||||
from docx.text.paragraph import Paragraph
|
||||
from docx.table import Table
|
||||
from docx.oxml.ns import qn
|
||||
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Please install python-docx to use the Word reader. "
|
||||
"You can install it by `pip install python-docx`.",
|
||||
) from e
|
||||
|
||||
doc = DocxDocument(word_path)
|
||||
|
||||
# If the last block is a table
|
||||
last_type = None
|
||||
|
||||
blocks: list[TextBlock | ImageBlock] = []
|
||||
for element in doc.element.body:
|
||||
if isinstance(element, CT_P):
|
||||
para = Paragraph(element, doc)
|
||||
|
||||
# Extract the text
|
||||
text = _extract_text_from_paragraph(para)
|
||||
|
||||
if self.include_image:
|
||||
# Check if the paragraph contains images
|
||||
has_drawing = bool(
|
||||
para._element.findall(".//" + qn("w:drawing")),
|
||||
)
|
||||
has_pict = bool(
|
||||
para._element.findall(".//" + qn("w:pict")),
|
||||
)
|
||||
|
||||
if has_drawing or has_pict:
|
||||
# Extract the image
|
||||
blocks.extend(_extract_image_data(para))
|
||||
last_type = "image"
|
||||
|
||||
# For current text block:
|
||||
# | separate_table | True | False |
|
||||
# |--------------------|--------|--------|
|
||||
# | last_type == text | append | append |
|
||||
# | last_type == image | new | new |
|
||||
# | last_type == table | new | append |
|
||||
# | last_type == None | new | new |
|
||||
if (
|
||||
last_type == "text"
|
||||
or last_type == "table"
|
||||
and not self.separate_table
|
||||
):
|
||||
blocks[-1]["text"] += "\n" + text
|
||||
else:
|
||||
blocks.append(
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=text,
|
||||
),
|
||||
)
|
||||
|
||||
# Update last type
|
||||
last_type = "text"
|
||||
|
||||
elif isinstance(element, CT_Tbl):
|
||||
# Extract the table data
|
||||
table_data = _extract_table_data(Table(element, doc))
|
||||
|
||||
if self.table_format == "markdown":
|
||||
text = self._table_to_markdown(table_data)
|
||||
else:
|
||||
text = self._table_to_json(table_data)
|
||||
|
||||
# For current table block:
|
||||
# | separate_table | True | False |
|
||||
# |--------------------|--------|--------|
|
||||
# | last_type == text | new | append |
|
||||
# | last_type == image | new | new |
|
||||
# | last_type == table | new | append |
|
||||
# | last_type == None | new | new |
|
||||
if not self.separate_table and last_type in ["text", "table"]:
|
||||
blocks[-1]["text"] += "\n" + text
|
||||
else:
|
||||
blocks.append(
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=text,
|
||||
),
|
||||
)
|
||||
|
||||
last_type = "table"
|
||||
|
||||
return blocks
|
||||
|
||||
@staticmethod
|
||||
def _table_to_markdown(table_data: list[list[str]]) -> str:
|
||||
"""Convert table data to Markdown format.
|
||||
|
||||
Args:
|
||||
table_data (`list[list[str]]`):
|
||||
Table data represented as a 2D list.
|
||||
|
||||
Returns:
|
||||
`str`:
|
||||
Table in Markdown format.
|
||||
"""
|
||||
if not table_data:
|
||||
return ""
|
||||
|
||||
num_cols = len(table_data[0])
|
||||
md_table = ""
|
||||
|
||||
# Header row
|
||||
header_row = "| " + " | ".join(table_data[0]) + " |\n"
|
||||
md_table += header_row
|
||||
|
||||
# Separator row
|
||||
separator_row = "| " + " | ".join(["---"] * num_cols) + " |\n"
|
||||
md_table += separator_row
|
||||
|
||||
# Data rows
|
||||
for row in table_data[1:]:
|
||||
data_row = "| " + " | ".join(row) + " |\n"
|
||||
md_table += data_row
|
||||
|
||||
return md_table
|
||||
|
||||
@staticmethod
|
||||
def _table_to_json(table_data: list[list[str]]) -> str:
|
||||
"""Convert table data to JSON string.
|
||||
|
||||
Args:
|
||||
table_data (`list[list[str]]`):
|
||||
Table data represented as a 2D list.
|
||||
|
||||
Returns:
|
||||
`str`:
|
||||
Table in JSON string format.
|
||||
"""
|
||||
json_strs = [
|
||||
"<system-info>A table loaded as a JSON array:</system-info>",
|
||||
]
|
||||
|
||||
for row in table_data:
|
||||
json_strs.append(
|
||||
json.dumps(row, ensure_ascii=False),
|
||||
)
|
||||
|
||||
return "\n".join(json_strs)
|
||||
|
||||
def get_doc_id(self, word_path: str) -> str:
|
||||
"""Generate a document ID based on the Word file path.
|
||||
|
||||
Args:
|
||||
word_path (`str`):
|
||||
The Word file path.
|
||||
|
||||
Returns:
|
||||
`str`:
|
||||
The generated document ID.
|
||||
"""
|
||||
return hashlib.md5(word_path.encode("utf-8")).hexdigest()
|
||||
@@ -0,0 +1,84 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""A general implementation of the knowledge class in AgentScope RAG module."""
|
||||
from typing import Any
|
||||
|
||||
from ._reader import Document
|
||||
from ..message import TextBlock
|
||||
from ._knowledge_base import KnowledgeBase
|
||||
|
||||
|
||||
class SimpleKnowledge(KnowledgeBase):
|
||||
"""A simple knowledge base implementation."""
|
||||
|
||||
async def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
limit: int = 5,
|
||||
score_threshold: float | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Document]:
|
||||
"""Retrieve relevant documents by the given queries.
|
||||
|
||||
Args:
|
||||
query (`str`):
|
||||
The query string to retrieve relevant documents.
|
||||
limit (`int`, defaults to 5):
|
||||
The number of relevant documents to retrieve.
|
||||
score_threshold: float | None = None,
|
||||
The threshold of the score to filter the results.
|
||||
**kwargs (`Any`):
|
||||
Other keyword arguments for the vector database search API.
|
||||
|
||||
Returns:
|
||||
`list[Document]`:
|
||||
A list of relevant documents.
|
||||
|
||||
TODO: handle the case when the query is too long.
|
||||
"""
|
||||
res_embedding = await self.embedding_model(
|
||||
[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=query,
|
||||
),
|
||||
],
|
||||
)
|
||||
res = await self.embedding_store.search(
|
||||
res_embedding.embeddings[0],
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
**kwargs,
|
||||
)
|
||||
return res
|
||||
|
||||
async def add_documents(
|
||||
self,
|
||||
documents: list[Document],
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Add documents to the knowledge
|
||||
|
||||
Args:
|
||||
documents (`list[Document]`):
|
||||
The list of documents to add.
|
||||
"""
|
||||
# Prepare the content to be embedded
|
||||
for doc in documents:
|
||||
if (
|
||||
doc.metadata.content["type"]
|
||||
not in self.embedding_model.supported_modalities
|
||||
):
|
||||
raise ValueError(
|
||||
f"The embedding model {self.embedding_model.model_name} "
|
||||
f"does not support {doc.metadata.content['type']} data.",
|
||||
)
|
||||
|
||||
# Get the embeddings
|
||||
res_embeddings = await self.embedding_model(
|
||||
[_.metadata.content for _ in documents],
|
||||
)
|
||||
|
||||
for doc, embedding in zip(documents, res_embeddings.embeddings):
|
||||
doc.embedding = embedding
|
||||
|
||||
await self.embedding_store.add(documents)
|
||||
@@ -0,0 +1,14 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The vector database store abstraction in AgentScope RAG module."""
|
||||
|
||||
from ._store_base import (
|
||||
VDBStoreBase,
|
||||
)
|
||||
from ._qdrant_store import QdrantStore
|
||||
from ._milvuslite_store import MilvusLiteStore
|
||||
|
||||
__all__ = [
|
||||
"VDBStoreBase",
|
||||
"QdrantStore",
|
||||
"MilvusLiteStore",
|
||||
]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,257 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The Milvus Lite vector store implementation."""
|
||||
import json
|
||||
from typing import Any, Literal, TYPE_CHECKING
|
||||
|
||||
from .._reader import Document
|
||||
from ._store_base import VDBStoreBase
|
||||
from .._document import DocMetadata
|
||||
|
||||
from ..._utils._common import _map_text_to_uuid
|
||||
from ...types import Embedding
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pymilvus import MilvusClient
|
||||
else:
|
||||
MilvusClient = "pymilvus.MilvusClient"
|
||||
|
||||
|
||||
class MilvusLiteStore(VDBStoreBase):
|
||||
"""The Milvus Lite vector store implementation, supporting both local and
|
||||
remote Milvus instances.
|
||||
|
||||
.. note:: In Milvus Lite, we use the scalar fields to store the metadata,
|
||||
including the document ID, chunk ID, and original content. The new
|
||||
MilvusClient API is used for simplified operations.
|
||||
|
||||
.. note:: Milvus Lite is not supported on Windows OS for now (2025-10-21).
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
uri: str,
|
||||
collection_name: str,
|
||||
dimensions: int,
|
||||
distance: Literal["COSINE", "L2", "IP"] = "COSINE",
|
||||
token: str = "",
|
||||
client_kwargs: dict[str, Any] | None = None,
|
||||
collection_kwargs: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Initialize the Milvus Lite vector store.
|
||||
|
||||
Args:
|
||||
uri (`str`):
|
||||
The URI of the Milvus instance. For Milvus Lite, use a local
|
||||
file path like "./milvus_demo.db". For remote Milvus server,
|
||||
use URI like "http://localhost:19530".
|
||||
collection_name (`str`):
|
||||
The name of the collection to store the embeddings.
|
||||
dimensions (`int`):
|
||||
The dimension of the embeddings.
|
||||
distance (`Literal["COSINE", "L2", "IP"]`, default to "COSINE"):
|
||||
The distance metric to use for the collection. Can be one of
|
||||
"COSINE", "L2", or "IP". Defaults to "COSINE".
|
||||
token (`str`, defaults to ""):
|
||||
The token for authentication when connecting to remote Milvus.
|
||||
Format: "username:password". Not needed for Milvus Lite.
|
||||
client_kwargs (`dict[str, Any] | None`, optional):
|
||||
Other keyword arguments for the Milvus client.
|
||||
collection_kwargs (`dict[str, Any] | None`, optional):
|
||||
Other keyword arguments for creating the collection.
|
||||
"""
|
||||
|
||||
try:
|
||||
from pymilvus import MilvusClient
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Milvus client is not installed. Please install it with "
|
||||
"`pip install pymilvus[milvus_lite]`.",
|
||||
) from e
|
||||
|
||||
client_kwargs = client_kwargs or {}
|
||||
|
||||
# Initialize MilvusClient with uri and optional token
|
||||
init_params = {"uri": uri, **client_kwargs}
|
||||
if token:
|
||||
init_params["token"] = token
|
||||
|
||||
self._client = MilvusClient(**init_params)
|
||||
|
||||
self.collection_name = collection_name
|
||||
self.dimensions = dimensions
|
||||
self.distance = distance
|
||||
self.collection_kwargs = collection_kwargs or {}
|
||||
|
||||
async def _validate_collection(self) -> None:
|
||||
"""Validate the collection exists, if not, create it."""
|
||||
if not self._client.has_collection(self.collection_name):
|
||||
# Create collection with the new MilvusClient API
|
||||
# By default, it creates an auto-incrementing integer ID field
|
||||
kwargs = {
|
||||
"collection_name": self.collection_name,
|
||||
"dimension": self.dimensions,
|
||||
"metric_type": self.distance,
|
||||
**self.collection_kwargs,
|
||||
}
|
||||
|
||||
self._client.create_collection(**kwargs)
|
||||
|
||||
async def add(self, documents: list[Document], **kwargs: Any) -> None:
|
||||
"""Add embeddings to the Milvus vector store.
|
||||
|
||||
Args:
|
||||
documents (`list[Document]`):
|
||||
A list of embedding records to be recorded in the Milvus store.
|
||||
**kwargs (`Any`):
|
||||
Additional arguments for the insert operation.
|
||||
"""
|
||||
await self._validate_collection()
|
||||
|
||||
# Prepare data for insertion using the new MilvusClient API
|
||||
data = []
|
||||
for doc in documents:
|
||||
# Generate a unique integer ID based on hash
|
||||
unique_string = json.dumps(
|
||||
{
|
||||
"doc_id": doc.metadata.doc_id,
|
||||
"chunk_id": doc.metadata.chunk_id,
|
||||
"content": doc.metadata.content,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
|
||||
id_type = self.collection_kwargs.get("id_type", "int")
|
||||
if id_type == "string":
|
||||
unique_id = _map_text_to_uuid(unique_string)[:6]
|
||||
else:
|
||||
unique_id = abs(hash(unique_string)) % (10**10)
|
||||
|
||||
# Prepare data entry with vector and metadata
|
||||
entry = {
|
||||
# Fixed fields for Milvus
|
||||
"id": unique_id,
|
||||
"vector": doc.embedding,
|
||||
# fields that will be returned in the "entity" field during
|
||||
# search
|
||||
"doc_id": doc.metadata.doc_id,
|
||||
"chunk_id": doc.metadata.chunk_id,
|
||||
"content": doc.metadata.content,
|
||||
"total_chunks": doc.metadata.total_chunks,
|
||||
}
|
||||
data.append(entry)
|
||||
|
||||
# Insert data using MilvusClient
|
||||
self._client.insert(
|
||||
collection_name=self.collection_name,
|
||||
data=data,
|
||||
)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query_embedding: Embedding,
|
||||
limit: int,
|
||||
score_threshold: float | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Document]:
|
||||
"""Search relevant documents from the Milvus vector store.
|
||||
|
||||
Args:
|
||||
query_embedding (`Embedding`):
|
||||
The embedding of the query text.
|
||||
limit (`int`):
|
||||
The number of relevant documents to retrieve.
|
||||
score_threshold (`float | None`, optional):
|
||||
The threshold of the score to filter the results.
|
||||
**kwargs (`Any`):
|
||||
Additional arguments for the Milvus client search API.
|
||||
- filter (`str`): Expression to filter the search results.
|
||||
- output_fields (`list[str]`): Fields to include in results.
|
||||
"""
|
||||
|
||||
# Get output fields if specified
|
||||
if "output_fields" not in kwargs:
|
||||
kwargs["output_fields"] = [
|
||||
"doc_id",
|
||||
"chunk_id",
|
||||
"content",
|
||||
"total_chunks",
|
||||
]
|
||||
|
||||
# Execute search using MilvusClient
|
||||
results = self._client.search(
|
||||
collection_name=self.collection_name,
|
||||
data=[query_embedding],
|
||||
limit=limit,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Process results
|
||||
collected_res = []
|
||||
for hits in results:
|
||||
for hit in hits:
|
||||
# Check score threshold
|
||||
if (
|
||||
score_threshold is not None
|
||||
and hit["distance"] < score_threshold
|
||||
):
|
||||
continue
|
||||
|
||||
# Get metadata from entity
|
||||
entity = hit["entity"]
|
||||
|
||||
doc_metadata = DocMetadata(
|
||||
content=entity.get("content", ""),
|
||||
doc_id=entity.get("doc_id", ""),
|
||||
chunk_id=entity.get("chunk_id", 0),
|
||||
total_chunks=entity.get("total_chunks", 0),
|
||||
)
|
||||
|
||||
# Create Document
|
||||
collected_res.append(
|
||||
Document(
|
||||
embedding=None, # Vector not returned by default
|
||||
score=hit["distance"],
|
||||
metadata=doc_metadata,
|
||||
),
|
||||
)
|
||||
|
||||
return collected_res
|
||||
|
||||
async def delete(
|
||||
self,
|
||||
ids: list[str] | None = None,
|
||||
filter: str | None = None, # pylint: disable=redefined-builtin
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Delete documents from the Milvus vector store.
|
||||
|
||||
Args:
|
||||
ids (`list[str] | None`, optional):
|
||||
List of entity IDs to delete.
|
||||
filter (`str | None`, optional):
|
||||
Expression to filter documents to delete.
|
||||
**kwargs (`Any`):
|
||||
Additional arguments for the delete operation.
|
||||
"""
|
||||
if ids is None and filter is None:
|
||||
raise ValueError(
|
||||
"Either ids or filter_expr must be provided for deletion.",
|
||||
)
|
||||
|
||||
# Delete data using MilvusClient
|
||||
self._client.delete(
|
||||
collection_name=self.collection_name,
|
||||
ids=ids,
|
||||
filter=filter,
|
||||
)
|
||||
|
||||
def get_client(self) -> MilvusClient:
|
||||
"""Get the underlying Milvus client, so that developers can access
|
||||
the full functionality of Milvus.
|
||||
|
||||
Returns:
|
||||
`MilvusClient`:
|
||||
The underlying Milvus client.
|
||||
"""
|
||||
return self._client
|
||||
@@ -0,0 +1,173 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The Qdrant local vector store implementation."""
|
||||
import json
|
||||
from typing import Any, Literal, TYPE_CHECKING
|
||||
|
||||
from .._reader import Document
|
||||
from ._store_base import VDBStoreBase
|
||||
from .._document import DocMetadata
|
||||
from ..._utils._common import _map_text_to_uuid
|
||||
from ...types import Embedding
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
else:
|
||||
AsyncQdrantClient = "qdrant_client.AsyncQdrantClient"
|
||||
|
||||
|
||||
class QdrantStore(VDBStoreBase):
|
||||
"""The Qdrant vector store implementation, supporting both local and
|
||||
remote Qdrant instances.
|
||||
|
||||
.. note:: In Qdrant, we use the ``payload`` field to store the metadata,
|
||||
including the document ID, chunk ID, and original content.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
location: Literal[":memory:"] | str,
|
||||
collection_name: str,
|
||||
dimensions: int,
|
||||
distance: Literal["Cosine", "Euclid", "Dot", "Manhattan"] = "Cosine",
|
||||
client_kwargs: dict[str, Any] | None = None,
|
||||
collection_kwargs: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Initialize the local Qdrant vector store.
|
||||
|
||||
Args:
|
||||
location (`Literal[":memory:"] | str`):
|
||||
The location of the Qdrant instance. Use ":memory:" for
|
||||
in-memory Qdrant instance, or url for remote Qdrant instance,
|
||||
e.g. "http://localhost:6333" or a path to a directory.
|
||||
collection_name (`str`):
|
||||
The name of the collection to store the embeddings.
|
||||
dimensions (`int`):
|
||||
The dimension of the embeddings.
|
||||
distance (`Literal["Cosine", "Euclid", "Dot", "Manhattan"]`, \
|
||||
default to "Cosine"):
|
||||
The distance metric to use for the collection. Can be one of
|
||||
"Cosine", "Euclid", "Dot", or "Manhattan". Defaults to
|
||||
"Cosine".
|
||||
client_kwargs (`dict[str, Any] | None`, optional):
|
||||
Other keyword arguments for the Qdrant client.
|
||||
collection_kwargs (`dict[str, Any] | None`, optional):
|
||||
Other keyword arguments for creating the collection.
|
||||
"""
|
||||
|
||||
try:
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Qdrant client is not installed. Please install it with "
|
||||
"`pip install qdrant-client`.",
|
||||
) from e
|
||||
|
||||
client_kwargs = client_kwargs or {}
|
||||
self._client = AsyncQdrantClient(location=location, **client_kwargs)
|
||||
|
||||
self.collection_name = collection_name
|
||||
self.dimensions = dimensions
|
||||
self.distance = distance
|
||||
self.collection_kwargs = collection_kwargs or {}
|
||||
|
||||
async def _validate_collection(self) -> None:
|
||||
"""Validate the collection exists, if not, create it."""
|
||||
if not await self._client.collection_exists(self.collection_name):
|
||||
from qdrant_client import models
|
||||
|
||||
collections_kwargs = {
|
||||
"collection_name": self.collection_name,
|
||||
"vectors_config": models.VectorParams(
|
||||
size=self.dimensions,
|
||||
distance=getattr(models.Distance, self.distance.upper()),
|
||||
),
|
||||
**self.collection_kwargs,
|
||||
}
|
||||
await self._client.create_collection(**collections_kwargs)
|
||||
|
||||
async def add(self, documents: list[Document], **kwargs: Any) -> None:
|
||||
"""Add embeddings to the Qdrant vector store.
|
||||
|
||||
Args:
|
||||
documents (`list[Document]`):
|
||||
A list of embedding records to be recorded in the Qdrant store.
|
||||
"""
|
||||
await self._validate_collection()
|
||||
|
||||
from qdrant_client.models import PointStruct
|
||||
|
||||
await self._client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=[
|
||||
PointStruct(
|
||||
id=_map_text_to_uuid(
|
||||
json.dumps(
|
||||
{
|
||||
"doc_id": _.metadata.doc_id,
|
||||
"chunk_id": _.metadata.chunk_id,
|
||||
"content": _.metadata.content,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
),
|
||||
vector=_.embedding,
|
||||
payload=_.metadata,
|
||||
)
|
||||
for _ in documents
|
||||
],
|
||||
)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query_embedding: Embedding,
|
||||
limit: int,
|
||||
score_threshold: float | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Document]:
|
||||
"""Search relevant documents from the Qdrant vector store.
|
||||
|
||||
Args:
|
||||
query_embedding (`Embedding`):
|
||||
The embedding of the query text.
|
||||
limit (`int`):
|
||||
The number of relevant documents to retrieve.
|
||||
score_threshold (`float | None`, optional):
|
||||
The threshold of the score to filter the results.
|
||||
**kwargs (`Any`):
|
||||
Other keyword arguments for the Qdrant client search API.
|
||||
"""
|
||||
res = await self._client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
query=query_embedding,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
collected_res = []
|
||||
for point in res.points:
|
||||
collected_res.append(
|
||||
Document(
|
||||
embedding=point.vector,
|
||||
score=point.score,
|
||||
metadata=DocMetadata(**point.payload),
|
||||
),
|
||||
)
|
||||
return collected_res
|
||||
|
||||
async def delete(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Delete is not implemented for QdrantStore."""
|
||||
raise NotImplementedError(
|
||||
"Delete is not implemented for QdrantStore.",
|
||||
)
|
||||
|
||||
def get_client(self) -> AsyncQdrantClient:
|
||||
"""Get the underlying Qdrant client, so that developers can access
|
||||
the full functionality of Qdrant.
|
||||
|
||||
Returns:
|
||||
`AsyncQdrantClient`:
|
||||
The underlying Qdrant client.
|
||||
"""
|
||||
return self._client
|
||||
@@ -0,0 +1,49 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The embedding store base class."""
|
||||
from abc import abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from .. import Document
|
||||
from ...types import Embedding
|
||||
|
||||
|
||||
class VDBStoreBase:
|
||||
"""The vector database store base class, serving as a middle layer between
|
||||
the knowledge base and the actual vector database implementation."""
|
||||
|
||||
@abstractmethod
|
||||
async def add(self, documents: list[Document], **kwargs: Any) -> None:
|
||||
"""Record the documents into the vector database."""
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Delete texts from the embedding store."""
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self,
|
||||
query_embedding: Embedding,
|
||||
limit: int,
|
||||
score_threshold: float | None = None,
|
||||
**kwargs: Any,
|
||||
) -> list[Document]:
|
||||
"""Retrieve relevant texts for the given queries.
|
||||
|
||||
Args:
|
||||
query_embedding (`Embedding`):
|
||||
The embedding of the query text.
|
||||
limit (`int`):
|
||||
The number of relevant documents to retrieve.
|
||||
score_threshold (`float | None`, optional):
|
||||
The threshold of the score to filter the results.
|
||||
**kwargs (`Any`):
|
||||
Other keyword arguments for the vector database search API.
|
||||
"""
|
||||
|
||||
def get_client(self) -> Any:
|
||||
"""Get the underlying vector database client, so that developers can
|
||||
access the full functionality of the vector database."""
|
||||
raise NotImplementedError(
|
||||
"``get_client`` is not implemented for "
|
||||
f"{self.__class__.__name__}.",
|
||||
)
|
||||
Reference in New Issue
Block a user