agentscope.rag

The retrieval-augmented generation (RAG) module in AgentScope.

class ReaderBase[source]

Bases: object

The reader base class, which is responsible for reading the original data, splitting it into chunks, and converting each chunk into a Document object.

abstract async __call__(*args, **kwargs)[source]

The async call function that takes the input files and returns the vector records

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

list[Document]

abstract get_doc_id(*args, **kwargs)[source]

Get a unique document ID for the input data. This method is to expose the document ID generation logic to the developers

Returns:

A unique document ID for the input data.

Return type:

str

Parameters:

  • args (Any)

  • kwargs (Any)

class TextReader[source]

Bases: ReaderBase

The text reader that splits text into chunks by a fixed chunk size and chunk overlap.

__init__(chunk_size=512, split_by='sentence')[source]

Initialize the text reader.

Parameters:

  • 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.

Return type:

None

async __call__(text)[source]

Read a text string, split it into chunks, and return a list of Document objects.

Parameters:

text (str) – The input text string, or a path to the local text file.

Returns:

A list of Document objects, where the metadata contains the chunked text, doc id and chunk id.

Return type:

list[Document]

get_doc_id(text)[source]

Get the document ID. This function can be used to check if the doc_id already exists in the knowledge base.

Parameters:

text (str)

Return type:

str

class PDFReader[source]

Bases: ReaderBase

The PDF reader that splits text into chunks by a fixed chunk size.

__init__(chunk_size=512, split_by='sentence')[source]

Initialize the text reader.

Parameters:

  • 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.

Return type:

None

async __call__(pdf_path)[source]

Read a PDF file, split it into chunks, and return a list of Document objects.

Parameters:

pdf_path (str) – The input PDF file path.

Return type:

list[Document]

get_doc_id(pdf_path)[source]

Get the document ID. This function can be used to check if the doc_id already exists in the knowledge base.

Parameters:

pdf_path (str)

Return type:

str

class ImageReader[source]

Bases: 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 __call__(image_url)[source]

Read an image and return the wrapped Document object.

Parameters:

image_url (str | list[str]) – The image URL(s) or path(s).

Returns:

A list of Document objects containing the image data.

Return type:

list[Document]

get_doc_id(image_path)[source]

Generate a document ID based on the image path.

Parameters:

image_path (str) – The image path or URL.

Returns:

The generated document ID.

Return type:

str

class DocMetadata[source]

Bases: 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.

__init__(content, doc_id, chunk_id, total_chunks)
Parameters:
Return type:

None

class Document[source]

Bases: object

The data chunk.

__init__(metadata, id=<factory>, embedding=<factory>, score=None)
Parameters:

  • metadata (DocMetadata)

  • id (str)

  • embedding (List[float] | None)

  • score (float | None)

Return type:

None

metadata: DocMetadata

The metadata of the data chunk.

id: str

The unique ID of the data chunk.

embedding: List[float] | None

The embedding of the data chunk.

score: float | None = None

The relevance score of the data chunk.

class VDBStoreBase[source]

Bases: object

The vector database store base class, serving as a middle layer between the knowledge base and the actual vector database implementation.

abstract async add(documents, **kwargs)[source]

Record the documents into the vector database.

Parameters:

  • documents (list[Document])

  • kwargs (Any)

Return type:

None

abstract async delete(*args, **kwargs)[source]

Delete texts from the embedding store.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

None

abstract async search(query_embedding, limit, score_threshold=None, **kwargs)[source]

Retrieve relevant texts for the given queries.

Parameters:

  • 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.

Return type:

list[Document]

get_client()[source]

Get the underlying vector database client, so that developers can access the full functionality of the vector database.

Return type:

Any

class QdrantStore[source]

Bases: 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.

__init__(location, collection_name, dimensions, distance='Cosine', client_kwargs=None, collection_kwargs=None)[source]

Initialize the local Qdrant vector store.

Parameters:

  • (`Literal[" (location) –

    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.

  • location (Literal[':memory:'] | str)

Return type:

None

async add(documents, **kwargs)[source]

Add embeddings to the Qdrant vector store.

Parameters:

  • documents (list[Document]) – A list of embedding records to be recorded in the Qdrant store.

  • kwargs (Any)

Return type:

None

async search(query_embedding, limit, score_threshold=None, **kwargs)[source]

Search relevant documents from the Qdrant vector store.

Parameters:

  • 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.

Return type:

list[Document]

async delete(*args, **kwargs)[source]

Delete is not implemented for QdrantStore.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

None

get_client()[source]

Get the underlying Qdrant client, so that developers can access the full functionality of Qdrant.

Returns:

The underlying Qdrant client.

Return type:

AsyncQdrantClient

class KnowledgeBase[source]

Bases: object

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.

__init__(embedding_store, embedding_model)[source]

Initialize the knowledge base.

Parameters:
Return type:

None

embedding_store: VDBStoreBase

The embedding store for the knowledge base.

embedding_model: EmbeddingModelBase

The embedding model for the knowledge base.

abstract async retrieve(query, limit=5, score_threshold=None, **kwargs)[source]

Retrieve relevant documents by the given query.

Parameters:

  • 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.

Return type:

list[Document]

abstract async add_documents(documents, **kwargs)[source]

Add documents to the knowledge base, which will embed the documents and store them in the embedding store.

Parameters:

  • documents (list[Document]) – A list of documents to add.

  • kwargs (Any)

Return type:

None

async retrieve_knowledge(query, limit=5, score_threshold=None, **kwargs)[source]

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.

Parameters:

  • 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.

  • kwargs (Any)

Return type:

ToolResponse

class SimpleKnowledge[source]

Bases: KnowledgeBase

A simple knowledge base implementation.

async retrieve(query, limit=5, score_threshold=None, **kwargs)[source]

Retrieve relevant documents by the given queries.

Parameters:

  • 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) – float | None = None, The threshold of the score to filter the results.

  • **kwargs (Any) – Other keyword arguments for the vector database search API.

Returns:

A list of relevant documents.

Return type:

list[Document]

TODO: handle the case when the query is too long.

async add_documents(documents, **kwargs)[source]

Add documents to the knowledge

Parameters:

  • documents (list[Document]) – The list of documents to add.

  • kwargs (Any)

Return type:

None