chore: 添加虚拟环境到仓库

- 添加 backend_service/venv 虚拟环境
- 包含所有Python依赖包
- 注意:虚拟环境约393MB,包含12655个文件
This commit is contained in:
2025-12-03 10:19:25 +08:00
parent a6c2027caa
commit c4f851d387
12655 changed files with 3009376 additions and 0 deletions

View File

@@ -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",
]

View File

@@ -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."""

View File

@@ -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.",
),
],
)

View File

@@ -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",
]

View File

@@ -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()

View File

@@ -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()

View File

@@ -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.
"""

View File

@@ -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()

View File

@@ -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()

View File

@@ -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)

View File

@@ -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",
]

View File

@@ -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

View File

@@ -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

View File

@@ -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__}.",
)