camel-ai 0.1.1__py3-none-any.whl → 0.1.4__py3-none-any.whl

camel/retrievers/vector_retriever.py

@@ -0,0 +1,183 @@
+# =========== Copyright 2023 @ 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 @ CAMEL-AI.org. All Rights Reserved. ===========
+from typing import Any, Dict, List, Optional
+
+from camel.embeddings import BaseEmbedding, OpenAIEmbedding
+from camel.loaders import UnstructuredIO
+from camel.retrievers.base import BaseRetriever
+from camel.storages import (
+    BaseVectorStorage,
+    QdrantStorage,
+    VectorDBQuery,
+    VectorRecord,
+)
+
+DEFAULT_TOP_K_RESULTS = 1
+DEFAULT_SIMILARITY_THRESHOLD = 0.75
+
+
+class VectorRetriever(BaseRetriever):
+    r"""An implementation of the `BaseRetriever` by using vector storage and
+    embedding model.
+
+    This class facilitates the retriever of relevant information using a
+    query-based approach, backed by vector embeddings.
+
+    Attributes:
+        embedding_model (BaseEmbedding): Embedding model used to generate
+            vector embeddings.
+        storage (BaseVectorStorage): Vector storage to query.
+        similarity_threshold (float, optional): The similarity threshold
+            for filtering results. Defaults to `DEFAULT_SIMILARITY_THRESHOLD`.
+        unstructured_modules (UnstructuredIO): A module for parsing files and
+            URLs and chunking content based on specified parameters.
+    """
+
+    def __init__(
+        self,
+        similarity_threshold: float = DEFAULT_SIMILARITY_THRESHOLD,
+        embedding_model: Optional[BaseEmbedding] = None,
+        storage: Optional[BaseVectorStorage] = None,
+    ) -> None:
+        r"""Initializes the retriever class with an optional embedding model.
+
+        Args:
+            similarity_threshold (float, optional): The similarity threshold
+                for filtering results. Defaults to
+                `DEFAULT_SIMILARITY_THRESHOLD`.
+            embedding_model (Optional[BaseEmbedding]): The embedding model
+                instance. Defaults to `OpenAIEmbedding` if not provided.
+            storage (BaseVectorStorage): Vector storage to query.
+        """
+        self.embedding_model = embedding_model or OpenAIEmbedding()
+        self.storage = (
+            storage
+            if storage is not None
+            else QdrantStorage(vector_dim=self.embedding_model.get_output_dim())
+        )
+        self.similarity_threshold = similarity_threshold
+        self.unstructured_modules: UnstructuredIO = UnstructuredIO()
+
+    def process(
+        self,
+        content_input_path: str,
+        chunk_type: str = "chunk_by_title",
+        **kwargs: Any,
+    ) -> None:
+        r"""Processes content from a file or URL, divides it into chunks by
+        using `Unstructured IO`, and stores their embeddings in the specified
+        vector storage.
+
+        Args:
+            content_input_path (str): File path or URL of the content to be
+                processed.
+            chunk_type (str): Type of chunking going to apply. Defaults to
+                "chunk_by_title".
+            **kwargs (Any): Additional keyword arguments for content parsing.
+        """
+        elements = self.unstructured_modules.parse_file_or_url(
+            content_input_path, **kwargs
+        )
+        chunks = self.unstructured_modules.chunk_elements(
+            chunk_type=chunk_type, elements=elements
+        )
+        # Iterate to process and store embeddings, set batch of 50
+        for i in range(0, len(chunks), 50):
+            batch_chunks = chunks[i : i + 50]
+            batch_vectors = self.embedding_model.embed_list(
+                objs=[str(chunk) for chunk in batch_chunks]
+            )
+
+            records = []
+            # Prepare the payload for each vector record, includes the content
+            # path, chunk metadata, and chunk text
+            for vector, chunk in zip(batch_vectors, batch_chunks):
+                content_path_info = {"content path": content_input_path}
+                chunk_metadata = {"metadata": chunk.metadata.to_dict()}
+                chunk_text = {"text": str(chunk)}
+                combined_dict = {
+                    **content_path_info,
+                    **chunk_metadata,
+                    **chunk_text,
+                }
+
+                records.append(
+                    VectorRecord(vector=vector, payload=combined_dict)
+                )
+
+            self.storage.add(records=records)
+
+    def query(
+        self,
+        query: str,
+        top_k: int = DEFAULT_TOP_K_RESULTS,
+    ) -> List[Dict[str, Any]]:
+        r"""Executes a query in vector storage and compiles the retrieved
+        results into a dictionary.
+
+        Args:
+            query (str): Query string for information retriever.
+            top_k (int, optional): The number of top results to return during
+                retriever. Must be a positive integer. Defaults to 1.
+
+        Returns:
+            List[Dict[str, Any]]: Concatenated list of the query results.
+
+        Raises:
+            ValueError: If 'top_k' is less than or equal to 0, if vector
+                storage is empty, if payload of vector storage is None.
+        """
+
+        if top_k <= 0:
+            raise ValueError("top_k must be a positive integer.")
+
+        # Load the storage incase it's hosted remote
+        self.storage.load()
+
+        query_vector = self.embedding_model.embed(obj=query)
+        db_query = VectorDBQuery(query_vector=query_vector, top_k=top_k)
+        query_results = self.storage.query(query=db_query)
+
+        if query_results[0].record.payload is None:
+            raise ValueError(
+                "Payload of vector storage is None, please check the collection."
+            )
+
+        # format the results
+        formatted_results = []
+        for result in query_results:
+            if (
+                result.similarity >= self.similarity_threshold
+                and result.record.payload is not None
+            ):
+                result_dict = {
+                    'similarity score': str(result.similarity),
+                    'content path': result.record.payload.get(
+                        'content path', ''
+                    ),
+                    'metadata': result.record.payload.get('metadata', {}),
+                    'text': result.record.payload.get('text', ''),
+                }
+                formatted_results.append(result_dict)
+
+        content_path = query_results[0].record.payload.get('content path', '')
+
+        if not formatted_results:
+            return [
+                {
+                    'text': f"""No suitable information retrieved from {content_path} \
+                with similarity_threshold = {self.similarity_threshold}."""
+                }
+            ]
+        return formatted_results

camel/societies/__init__.py

@@ -11,8 +11,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
-from .role_playing import RolePlaying
 from .babyagi_playing import BabyAGI
+from .role_playing import RolePlaying
 
 __all__ = [
     'RolePlaying',

camel/societies/babyagi_playing.py

@@ -80,13 +80,17 @@ class BabyAGI:
         self.task_type = task_type
         self.task_prompt = task_prompt
         self.specified_task_prompt: TextPrompt
-        self.init_specified_task_prompt(assistant_role_name, user_role_name,
-                                        task_specify_agent_kwargs,
-                                        extend_task_specify_meta_dict,
-                                        output_language)
+        self.init_specified_task_prompt(
+            assistant_role_name,
+            user_role_name,
+            task_specify_agent_kwargs,
+            extend_task_specify_meta_dict,
+            output_language,
+        )
 
         sys_msg_generator = SystemMessageGenerator(
-            task_type=self.task_type, **(sys_msg_generator_kwargs or {}))
+            task_type=self.task_type, **(sys_msg_generator_kwargs or {})
+        )
 
         init_assistant_sys_msg = sys_msg_generator.from_dicts(
             meta_dicts=[
@@ -105,20 +109,27 @@ class BabyAGI:
         self.assistant_sys_msg: BaseMessage
         self.task_creation_agent: TaskCreationAgent
         self.task_prioritization_agent: TaskPrioritizationAgent
-        self.init_agents(init_assistant_sys_msg[0], assistant_agent_kwargs,
-                         task_creation_agent_kwargs,
-                         task_prioritization_agent_kwargs, output_language,
-                         message_window_size)
+        self.init_agents(
+            init_assistant_sys_msg[0],
+            assistant_agent_kwargs,
+            task_creation_agent_kwargs,
+            task_prioritization_agent_kwargs,
+            output_language,
+            message_window_size,
+        )
 
         self.subtasks: deque = deque([])
         self.solved_subtasks: List[str] = []
         self.MAX_TASK_HISTORY = max_task_history
 
     def init_specified_task_prompt(
-            self, assistant_role_name: str, user_role_name: str,
-            task_specify_agent_kwargs: Optional[Dict],
-            extend_task_specify_meta_dict: Optional[Dict],
-            output_language: Optional[str]):
+        self,
+        assistant_role_name: str,
+        user_role_name: str,
+        task_specify_agent_kwargs: Optional[Dict],
+        extend_task_specify_meta_dict: Optional[Dict],
+        output_language: Optional[str],
+    ):
         r"""Use a task specify agent to generate a specified task prompt.
         Generated specified task prompt will be used to replace original
         task prompt. If there is no task specify agent, specified task
@@ -138,8 +149,10 @@ class BabyAGI:
         task_specify_meta_dict = dict()
         if self.task_type in [TaskType.AI_SOCIETY, TaskType.MISALIGNMENT]:
             task_specify_meta_dict.update(
-                dict(assistant_role=assistant_role_name,
-                     user_role=user_role_name))
+                dict(
+                    assistant_role=assistant_role_name, user_role=user_role_name
+                )
+            )
         task_specify_meta_dict.update(extend_task_specify_meta_dict or {})
         task_specify_agent = TaskSpecifyAgent(
             task_type=self.task_type,
@@ -151,12 +164,15 @@ class BabyAGI:
             meta_dict=task_specify_meta_dict,
         )
 
-    def init_agents(self, init_assistant_sys_msg: BaseMessage,
-                    assistant_agent_kwargs: Optional[Dict],
-                    task_creation_agent_kwargs: Optional[Dict],
-                    task_prioritization_agent_kwargs: Optional[Dict],
-                    output_language: Optional[str],
-                    message_window_size: Optional[int] = None):
+    def init_agents(
+        self,
+        init_assistant_sys_msg: BaseMessage,
+        assistant_agent_kwargs: Optional[Dict],
+        task_creation_agent_kwargs: Optional[Dict],
+        task_prioritization_agent_kwargs: Optional[Dict],
+        output_language: Optional[str],
+        message_window_size: Optional[int] = None,
+    ):
         r"""Initialize assistant and user agents with their system messages.
 
         Args:
@@ -215,12 +231,14 @@ class BabyAGI:
         if not self.subtasks:
             new_subtask_list = self.task_creation_agent.run(task_list=[])
             prioritized_subtask_list = self.task_prioritization_agent.run(
-                new_subtask_list)
+                new_subtask_list
+            )
             self.subtasks = deque(prioritized_subtask_list)
 
         task_name = self.subtasks.popleft()
         assistant_msg_msg = BaseMessage.make_user_message(
-            role_name=self.assistant_sys_msg.role_name, content=f"{task_name}")
+            role_name=self.assistant_sys_msg.role_name, content=f"{task_name}"
+        )
 
         assistant_response = self.assistant_agent.step(assistant_msg_msg)
         assistant_msg = assistant_response.msgs[0]
@@ -232,12 +250,14 @@ class BabyAGI:
         past_tasks = self.solved_subtasks + list(self.subtasks)
 
         new_subtask_list = self.task_creation_agent.run(
-            task_list=past_tasks[-self.MAX_TASK_HISTORY:])
+            task_list=past_tasks[-self.MAX_TASK_HISTORY :]
+        )
 
         if new_subtask_list:
             self.subtasks.extend(new_subtask_list)
             prioritized_subtask_list = self.task_prioritization_agent.run(
-                task_list=list(self.subtasks)[-self.MAX_TASK_HISTORY:])
+                task_list=list(self.subtasks)[-self.MAX_TASK_HISTORY :]
+            )
             self.subtasks = deque(prioritized_subtask_list)
         else:
             print("no new tasks")
@@ -245,10 +265,14 @@ class BabyAGI:
         assistant_response.info['subtasks'] = list(self.subtasks)
         if not self.subtasks:
             terminated = True
-            assistant_response.info[
-                'termination_reasons'] = "All tasks are solved"
-            return ChatAgentResponse([assistant_msg], terminated,
-                                     assistant_response.info)
-        return ChatAgentResponse([assistant_msg],
-                                 assistant_response.terminated,
-                                 assistant_response.info)
+            assistant_response.info['termination_reasons'] = (
+                "All tasks are solved"
+            )
+            return ChatAgentResponse(
+                [assistant_msg], terminated, assistant_response.info
+            )
+        return ChatAgentResponse(
+            [assistant_msg],
+            assistant_response.terminated,
+            assistant_response.info,
+        )