File size: 3,879 Bytes
fcaa164 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
from typing import List, Optional
from camel.embeddings import BaseEmbedding, OpenAIEmbedding
from camel.memories.base import MemoryBlock
from camel.memories.records import ContextRecord, MemoryRecord
from camel.storages.vectordb_storages import (
BaseVectorStorage,
QdrantStorage,
VectorDBQuery,
VectorRecord,
)
class VectorDBBlock(MemoryBlock):
r"""An implementation of the :obj:`MemoryBlock` abstract base class for
maintaining and retrieving information using vector embeddings within a
vector database.
Args:
storage (Optional[BaseVectorStorage], optional): The storage mechanism
for the vector database. Defaults to in-memory :obj:`Qdrant` if not
provided. (default: :obj:`None`)
embedding (Optional[BaseEmbedding], optional): Embedding mechanism to
convert chat messages into vector representations. Defaults to
:obj:`OpenAiEmbedding` if not provided. (default: :obj:`None`)
"""
def __init__(
self,
storage: Optional[BaseVectorStorage] = None,
embedding: Optional[BaseEmbedding] = None,
) -> None:
self.embedding = embedding or OpenAIEmbedding()
self.vector_dim = self.embedding.get_output_dim()
self.storage = storage or QdrantStorage(vector_dim=self.vector_dim)
def retrieve(
self,
keyword: str,
limit: int = 3,
) -> List[ContextRecord]:
r"""Retrieves similar records from the vector database based on the
content of the keyword.
Args:
keyword (str): This string will be converted into a vector
representation to query the database.
limit (int, optional): The maximum number of similar messages to
retrieve. (default: :obj:`3`).
Returns:
List[ContextRecord]: A list of memory records retrieved from the
vector database based on similarity to :obj:`current_state`.
"""
query_vector = self.embedding.embed(keyword)
results = self.storage.query(
VectorDBQuery(query_vector=query_vector, top_k=limit)
)
return [
ContextRecord(
memory_record=MemoryRecord.from_dict(result.record.payload),
score=result.similarity,
)
for result in results
if result.record.payload is not None
]
def write_records(self, records: List[MemoryRecord]) -> None:
"""
Converts the provided chat messages into vector representations and
writes them to the vector database.
Args:
records (List[MemoryRecord]): Memory records to be added to the
memory.
"""
v_records = [
VectorRecord(
vector=self.embedding.embed(record.message.content),
payload=record.to_dict(),
id=str(record.uuid),
)
for record in records
]
self.storage.add(v_records)
def clear(self) -> None:
r"""Removes all records from the vector database memory."""
self.storage.clear()
|