ag2 0.4b1__py3-none-any.whl → 0.4.2b1__py3-none-any.whl

autogen/agentchat/contrib/tool_retriever.py

@@ -1,114 +0,0 @@
-import importlib.util
-import inspect
-import os
-from textwrap import dedent, indent
-
-import pandas as pd
-from sentence_transformers import SentenceTransformer, util
-
-from autogen import AssistantAgent, UserProxyAgent
-from autogen.coding import LocalCommandLineCodeExecutor
-
-
-class ToolBuilder:
-    TOOL_USING_PROMPT = """# Functions
-    You have access to the following functions. They can be accessed from the module called 'functions' by their function names.
-For example, if there is a function called `foo` you could import it by writing `from functions import foo`
-{functions}
-"""
-
-    def __init__(self, corpus_path, retriever="all-mpnet-base-v2"):
-
-        self.df = pd.read_csv(corpus_path, sep="\t")
-        document_list = self.df["document_content"].tolist()
-
-        self.model = SentenceTransformer(retriever)
-        self.embeddings = self.model.encode(document_list)
-
-    def retrieve(self, query, top_k=3):
-        # Encode the query using the Sentence Transformer model
-        query_embedding = self.model.encode([query])
-
-        hits = util.semantic_search(query_embedding, self.embeddings, top_k=top_k)
-
-        results = []
-        for hit in hits[0]:
-            results.append(self.df.iloc[hit["corpus_id"], 1])
-        return results
-
-    def bind(self, agent: AssistantAgent, functions: str):
-        """Binds the function to the agent so that agent is aware of it."""
-        sys_message = agent.system_message
-        sys_message += self.TOOL_USING_PROMPT.format(functions=functions)
-        agent.update_system_message(sys_message)
-        return
-
-    def bind_user_proxy(self, agent: UserProxyAgent, tool_root: str):
-        """
-        Updates user proxy agent with a executor so that code executor can successfully execute function-related code.
-        Returns an updated user proxy.
-        """
-        # Find all the functions in the tool root
-        functions = find_callables(tool_root)
-
-        code_execution_config = agent._code_execution_config
-        executor = LocalCommandLineCodeExecutor(
-            timeout=code_execution_config.get("timeout", 180),
-            work_dir=code_execution_config.get("work_dir", "coding"),
-            functions=functions,
-        )
-        code_execution_config = {
-            "executor": executor,
-            "last_n_messages": code_execution_config.get("last_n_messages", 1),
-        }
-        updated_user_proxy = UserProxyAgent(
-            name=agent.name,
-            is_termination_msg=agent._is_termination_msg,
-            code_execution_config=code_execution_config,
-            human_input_mode="NEVER",
-            default_auto_reply=agent._default_auto_reply,
-        )
-        return updated_user_proxy
-
-
-def get_full_tool_description(py_file):
-    """
-    Retrieves the function signature for a given Python file.
-    """
-    with open(py_file, "r") as f:
-        code = f.read()
-        exec(code)
-        function_name = os.path.splitext(os.path.basename(py_file))[0]
-        if function_name in locals():
-            func = locals()[function_name]
-            content = f"def {func.__name__}{inspect.signature(func)}:\n"
-            docstring = func.__doc__
-
-            if docstring:
-                docstring = dedent(docstring)
-                docstring = '"""' + docstring + '"""'
-                docstring = indent(docstring, "    ")
-                content += docstring + "\n"
-            return content
-        else:
-            raise ValueError(f"Function {function_name} not found in {py_file}")
-
-
-def find_callables(directory):
-    """
-    Find all callable objects defined in Python files within the specified directory.
-    """
-    callables = []
-    for root, dirs, files in os.walk(directory):
-        for file in files:
-            if file.endswith(".py"):
-                module_name = os.path.splitext(file)[0]
-                module_path = os.path.join(root, file)
-                spec = importlib.util.spec_from_file_location(module_name, module_path)
-                module = importlib.util.module_from_spec(spec)
-                spec.loader.exec_module(module)
-                for name, value in module.__dict__.items():
-                    if callable(value) and name == module_name:
-                        callables.append(value)
-                        break
-    return callables

autogen/agentchat/contrib/vectordb/base.py

@@ -1,243 +0,0 @@
-# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
-#
-# SPDX-License-Identifier: Apache-2.0
-#
-# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
-# SPDX-License-Identifier: MIT
-from typing import (
-    Any,
-    Callable,
-    List,
-    Mapping,
-    Optional,
-    Protocol,
-    Sequence,
-    Tuple,
-    TypedDict,
-    Union,
-    runtime_checkable,
-)
-
-Metadata = Union[Mapping[str, Any], None]
-Vector = Union[Sequence[float], Sequence[int]]
-ItemID = Union[str, int]  # chromadb doesn't support int ids, VikingDB does
-
-
-class Document(TypedDict):
-    """A Document is a record in the vector database.
-
-    id: ItemID | the unique identifier of the document.
-    content: str | the text content of the chunk.
-    metadata: Metadata, Optional | contains additional information about the document such as source, date, etc.
-    embedding: Vector, Optional | the vector representation of the content.
-    """
-
-    id: ItemID
-    content: str
-    metadata: Optional[Metadata]
-    embedding: Optional[Vector]
-
-
-"""QueryResults is the response from the vector database for a query/queries.
-A query is a list containing one string while queries is a list containing multiple strings.
-The response is a list of query results, each query result is a list of tuples containing the document and the distance.
-"""
-QueryResults = List[List[Tuple[Document, float]]]
-
-
-@runtime_checkable
-class VectorDB(Protocol):
-    """
-    Abstract class for vector database. A vector database is responsible for storing and retrieving documents.
-
-    Attributes:
-        active_collection: Any | The active collection in the vector database. Make get_collection faster. Default is None.
-        type: str | The type of the vector database, chroma, pgvector, etc. Default is "".
-
-    Methods:
-        create_collection: Callable[[str, bool, bool], Any] | Create a collection in the vector database.
-        get_collection: Callable[[str], Any] | Get the collection from the vector database.
-        delete_collection: Callable[[str], Any] | Delete the collection from the vector database.
-        insert_docs: Callable[[List[Document], str, bool], None] | Insert documents into the collection of the vector database.
-        update_docs: Callable[[List[Document], str], None] | Update documents in the collection of the vector database.
-        delete_docs: Callable[[List[ItemID], str], None] | Delete documents from the collection of the vector database.
-        retrieve_docs: Callable[[List[str], str, int, float], QueryResults] | Retrieve documents from the collection of the vector database based on the queries.
-        get_docs_by_ids: Callable[[List[ItemID], str], List[Document]] | Retrieve documents from the collection of the vector database based on the ids.
-    """
-
-    active_collection: Any = None
-    type: str = ""
-    embedding_function: Optional[Callable[[List[str]], List[List[float]]]] = (
-        None  # embeddings = embedding_function(sentences)
-    )
-
-    def create_collection(self, collection_name: str, overwrite: bool = False, get_or_create: bool = True) -> Any:
-        """
-        Create a collection in the vector database.
-        Case 1. if the collection does not exist, create the collection.
-        Case 2. the collection exists, if overwrite is True, it will overwrite the collection.
-        Case 3. the collection exists and overwrite is False, if get_or_create is True, it will get the collection,
-            otherwise it raise a ValueError.
-
-        Args:
-            collection_name: str | The name of the collection.
-            overwrite: bool | Whether to overwrite the collection if it exists. Default is False.
-            get_or_create: bool | Whether to get the collection if it exists. Default is True.
-
-        Returns:
-            Any | The collection object.
-        """
-        ...
-
-    def get_collection(self, collection_name: str = None) -> Any:
-        """
-        Get the collection from the vector database.
-
-        Args:
-            collection_name: str | The name of the collection. Default is None. If None, return the
-                current active collection.
-
-        Returns:
-            Any | The collection object.
-        """
-        ...
-
-    def delete_collection(self, collection_name: str) -> Any:
-        """
-        Delete the collection from the vector database.
-
-        Args:
-            collection_name: str | The name of the collection.
-
-        Returns:
-            Any
-        """
-        ...
-
-    def insert_docs(self, docs: List[Document], collection_name: str = None, upsert: bool = False, **kwargs) -> None:
-        """
-        Insert documents into the collection of the vector database.
-
-        Args:
-            docs: List[Document] | A list of documents. Each document is a TypedDict `Document`.
-            collection_name: str | The name of the collection. Default is None.
-            upsert: bool | Whether to update the document if it exists. Default is False.
-            kwargs: Dict | Additional keyword arguments.
-
-        Returns:
-            None
-        """
-        ...
-
-    def update_docs(self, docs: List[Document], collection_name: str = None, **kwargs) -> None:
-        """
-        Update documents in the collection of the vector database.
-
-        Args:
-            docs: List[Document] | A list of documents.
-            collection_name: str | The name of the collection. Default is None.
-            kwargs: Dict | Additional keyword arguments.
-
-        Returns:
-            None
-        """
-        ...
-
-    def delete_docs(self, ids: List[ItemID], collection_name: str = None, **kwargs) -> None:
-        """
-        Delete documents from the collection of the vector database.
-
-        Args:
-            ids: List[ItemID] | A list of document ids. Each id is a typed `ItemID`.
-            collection_name: str | The name of the collection. Default is None.
-            kwargs: Dict | Additional keyword arguments.
-
-        Returns:
-            None
-        """
-        ...
-
-    def retrieve_docs(
-        self,
-        queries: List[str],
-        collection_name: str = None,
-        n_results: int = 10,
-        distance_threshold: float = -1,
-        **kwargs,
-    ) -> QueryResults:
-        """
-        Retrieve documents from the collection of the vector database based on the queries.
-
-        Args:
-            queries: List[str] | A list of queries. Each query is a string.
-            collection_name: str | The name of the collection. Default is None.
-            n_results: int | The number of relevant documents to return. Default is 10.
-            distance_threshold: float | The threshold for the distance score, only distance smaller than it will be
-                returned. Don't filter with it if < 0. Default is -1.
-            kwargs: Dict | Additional keyword arguments.
-
-        Returns:
-            QueryResults | The query results. Each query result is a list of list of tuples containing the document and
-                the distance.
-        """
-        ...
-
-    def get_docs_by_ids(
-        self, ids: List[ItemID] = None, collection_name: str = None, include=None, **kwargs
-    ) -> List[Document]:
-        """
-        Retrieve documents from the collection of the vector database based on the ids.
-
-        Args:
-            ids: List[ItemID] | A list of document ids. If None, will return all the documents. Default is None.
-            collection_name: str | The name of the collection. Default is None.
-            include: List[str] | The fields to include. Default is None.
-                If None, will include ["metadatas", "documents"], ids will always be included. This may differ
-                depending on the implementation.
-            kwargs: dict | Additional keyword arguments.
-
-        Returns:
-            List[Document] | The results.
-        """
-        ...
-
-
-class VectorDBFactory:
-    """
-    Factory class for creating vector databases.
-    """
-
-    PREDEFINED_VECTOR_DB = ["chroma", "pgvector", "mongodb", "qdrant"]
-
-    @staticmethod
-    def create_vector_db(db_type: str, **kwargs) -> VectorDB:
-        """
-        Create a vector database.
-
-        Args:
-            db_type: str | The type of the vector database.
-            kwargs: Dict | The keyword arguments for initializing the vector database.
-
-        Returns:
-            VectorDB | The vector database.
-        """
-        if db_type.lower() in ["chroma", "chromadb"]:
-            from .chromadb import ChromaVectorDB
-
-            return ChromaVectorDB(**kwargs)
-        if db_type.lower() in ["pgvector", "pgvectordb"]:
-            from .pgvectordb import PGVectorDB
-
-            return PGVectorDB(**kwargs)
-        if db_type.lower() in ["mdb", "mongodb", "atlas"]:
-            from .mongodb import MongoDBAtlasVectorDB
-
-            return MongoDBAtlasVectorDB(**kwargs)
-        if db_type.lower() in ["qdrant", "qdrantdb"]:
-            from .qdrant import QdrantVectorDB
-
-            return QdrantVectorDB(**kwargs)
-        else:
-            raise ValueError(
-                f"Unsupported vector database type: {db_type}. Valid types are {VectorDBFactory.PREDEFINED_VECTOR_DB}."
-            )

autogen/agentchat/contrib/vectordb/chromadb.py

@@ -1,326 +0,0 @@
-# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
-#
-# SPDX-License-Identifier: Apache-2.0
-#
-# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
-# SPDX-License-Identifier: MIT
-import os
-from typing import Callable, List
-
-from .base import Document, ItemID, QueryResults, VectorDB
-from .utils import chroma_results_to_query_results, filter_results_by_distance, get_logger
-
-try:
-    import chromadb
-
-    if chromadb.__version__ < "0.4.15":
-        raise ImportError("Please upgrade chromadb to version 0.4.15 or later.")
-    import chromadb.utils.embedding_functions as ef
-    from chromadb.api.models.Collection import Collection
-except ImportError:
-    raise ImportError("Please install chromadb: `pip install chromadb`")
-
-CHROMADB_MAX_BATCH_SIZE = os.environ.get("CHROMADB_MAX_BATCH_SIZE", 40000)
-logger = get_logger(__name__)
-
-
-class ChromaVectorDB(VectorDB):
-    """
-    A vector database that uses ChromaDB as the backend.
-    """
-
-    def __init__(
-        self, *, client=None, path: str = "tmp/db", embedding_function: Callable = None, metadata: dict = None, **kwargs
-    ) -> None:
-        """
-        Initialize the vector database.
-
-        Args:
-            client: chromadb.Client | The client object of the vector database. Default is None.
-                If provided, it will use the client object directly and ignore other arguments.
-            path: str | The path to the vector database. Default is `tmp/db`. The default was `None` for version <=0.2.24.
-            embedding_function: Callable | The embedding function used to generate the vector representation
-                of the documents. Default is None, SentenceTransformerEmbeddingFunction("all-MiniLM-L6-v2") will be used.
-            metadata: dict | The metadata of the vector database. Default is None. If None, it will use this
-                setting: {"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32}. For more details of
-                the metadata, please refer to [distances](https://github.com/nmslib/hnswlib#supported-distances),
-                [hnsw](https://github.com/chroma-core/chroma/blob/566bc80f6c8ee29f7d99b6322654f32183c368c4/chromadb/segment/impl/vector/local_hnsw.py#L184),
-                and [ALGO_PARAMS](https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md).
-            kwargs: dict | Additional keyword arguments.
-
-        Returns:
-            None
-        """
-        self.client = client
-        self.path = path
-        self.embedding_function = (
-            ef.SentenceTransformerEmbeddingFunction("all-MiniLM-L6-v2")
-            if embedding_function is None
-            else embedding_function
-        )
-        self.metadata = metadata if metadata else {"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32}
-        if not self.client:
-            if self.path is not None:
-                self.client = chromadb.PersistentClient(path=self.path, **kwargs)
-            else:
-                self.client = chromadb.Client(**kwargs)
-        self.active_collection = None
-        self.type = "chroma"
-
-    def create_collection(
-        self, collection_name: str, overwrite: bool = False, get_or_create: bool = True
-    ) -> Collection:
-        """
-        Create a collection in the vector database.
-        Case 1. if the collection does not exist, create the collection.
-        Case 2. the collection exists, if overwrite is True, it will overwrite the collection.
-        Case 3. the collection exists and overwrite is False, if get_or_create is True, it will get the collection,
-            otherwise it raise a ValueError.
-
-        Args:
-            collection_name: str | The name of the collection.
-            overwrite: bool | Whether to overwrite the collection if it exists. Default is False.
-            get_or_create: bool | Whether to get the collection if it exists. Default is True.
-
-        Returns:
-            Collection | The collection object.
-        """
-        try:
-            if self.active_collection and self.active_collection.name == collection_name:
-                collection = self.active_collection
-            else:
-                collection = self.client.get_collection(collection_name, embedding_function=self.embedding_function)
-        except ValueError:
-            collection = None
-        if collection is None:
-            return self.client.create_collection(
-                collection_name,
-                embedding_function=self.embedding_function,
-                get_or_create=get_or_create,
-                metadata=self.metadata,
-            )
-        elif overwrite:
-            self.client.delete_collection(collection_name)
-            return self.client.create_collection(
-                collection_name,
-                embedding_function=self.embedding_function,
-                get_or_create=get_or_create,
-                metadata=self.metadata,
-            )
-        elif get_or_create:
-            return collection
-        else:
-            raise ValueError(f"Collection {collection_name} already exists.")
-
-    def get_collection(self, collection_name: str = None) -> Collection:
-        """
-        Get the collection from the vector database.
-
-        Args:
-            collection_name: str | The name of the collection. Default is None. If None, return the
-                current active collection.
-
-        Returns:
-            Collection | The collection object.
-        """
-        if collection_name is None:
-            if self.active_collection is None:
-                raise ValueError("No collection is specified.")
-            else:
-                logger.info(
-                    f"No collection is specified. Using current active collection {self.active_collection.name}."
-                )
-        else:
-            if not (self.active_collection and self.active_collection.name == collection_name):
-                self.active_collection = self.client.get_collection(
-                    collection_name, embedding_function=self.embedding_function
-                )
-        return self.active_collection
-
-    def delete_collection(self, collection_name: str) -> None:
-        """
-        Delete the collection from the vector database.
-
-        Args:
-            collection_name: str | The name of the collection.
-
-        Returns:
-            None
-        """
-        self.client.delete_collection(collection_name)
-        if self.active_collection and self.active_collection.name == collection_name:
-            self.active_collection = None
-
-    def _batch_insert(
-        self, collection: Collection, embeddings=None, ids=None, metadatas=None, documents=None, upsert=False
-    ) -> None:
-        batch_size = int(CHROMADB_MAX_BATCH_SIZE)
-        for i in range(0, len(documents), min(batch_size, len(documents))):
-            end_idx = i + min(batch_size, len(documents) - i)
-            collection_kwargs = {
-                "documents": documents[i:end_idx],
-                "ids": ids[i:end_idx],
-                "metadatas": metadatas[i:end_idx] if metadatas else None,
-                "embeddings": embeddings[i:end_idx] if embeddings else None,
-            }
-            if upsert:
-                collection.upsert(**collection_kwargs)
-            else:
-                collection.add(**collection_kwargs)
-
-    def insert_docs(self, docs: List[Document], collection_name: str = None, upsert: bool = False) -> None:
-        """
-        Insert documents into the collection of the vector database.
-
-        Args:
-            docs: List[Document] | A list of documents. Each document is a TypedDict `Document`.
-            collection_name: str | The name of the collection. Default is None.
-            upsert: bool | Whether to update the document if it exists. Default is False.
-            kwargs: Dict | Additional keyword arguments.
-
-        Returns:
-            None
-        """
-        if not docs:
-            return
-        if docs[0].get("content") is None:
-            raise ValueError("The document content is required.")
-        if docs[0].get("id") is None:
-            raise ValueError("The document id is required.")
-        documents = [doc.get("content") for doc in docs]
-        ids = [doc.get("id") for doc in docs]
-        collection = self.get_collection(collection_name)
-        if docs[0].get("embedding") is None:
-            logger.info(
-                "No content embedding is provided. Will use the VectorDB's embedding function to generate the content embedding."
-            )
-            embeddings = None
-        else:
-            embeddings = [doc.get("embedding") for doc in docs]
-        if docs[0].get("metadata") is None:
-            metadatas = None
-        else:
-            metadatas = [doc.get("metadata") for doc in docs]
-        self._batch_insert(collection, embeddings, ids, metadatas, documents, upsert)
-
-    def update_docs(self, docs: List[Document], collection_name: str = None) -> None:
-        """
-        Update documents in the collection of the vector database.
-
-        Args:
-            docs: List[Document] | A list of documents.
-            collection_name: str | The name of the collection. Default is None.
-
-        Returns:
-            None
-        """
-        self.insert_docs(docs, collection_name, upsert=True)
-
-    def delete_docs(self, ids: List[ItemID], collection_name: str = None, **kwargs) -> None:
-        """
-        Delete documents from the collection of the vector database.
-
-        Args:
-            ids: List[ItemID] | A list of document ids. Each id is a typed `ItemID`.
-            collection_name: str | The name of the collection. Default is None.
-            kwargs: Dict | Additional keyword arguments.
-
-        Returns:
-            None
-        """
-        collection = self.get_collection(collection_name)
-        collection.delete(ids, **kwargs)
-
-    def retrieve_docs(
-        self,
-        queries: List[str],
-        collection_name: str = None,
-        n_results: int = 10,
-        distance_threshold: float = -1,
-        **kwargs,
-    ) -> QueryResults:
-        """
-        Retrieve documents from the collection of the vector database based on the queries.
-
-        Args:
-            queries: List[str] | A list of queries. Each query is a string.
-            collection_name: str | The name of the collection. Default is None.
-            n_results: int | The number of relevant documents to return. Default is 10.
-            distance_threshold: float | The threshold for the distance score, only distance smaller than it will be
-                returned. Don't filter with it if < 0. Default is -1.
-            kwargs: Dict | Additional keyword arguments.
-
-        Returns:
-            QueryResults | The query results. Each query result is a list of list of tuples containing the document and
-                the distance.
-        """
-        collection = self.get_collection(collection_name)
-        if isinstance(queries, str):
-            queries = [queries]
-        results = collection.query(
-            query_texts=queries,
-            n_results=n_results,
-            **kwargs,
-        )
-        results["contents"] = results.pop("documents")
-        results = chroma_results_to_query_results(results)
-        results = filter_results_by_distance(results, distance_threshold)
-        return results
-
-    @staticmethod
-    def _chroma_get_results_to_list_documents(data_dict) -> List[Document]:
-        """Converts a dictionary with list values to a list of Document.
-
-        Args:
-            data_dict: A dictionary where keys map to lists or None.
-
-        Returns:
-            List[Document] | The list of Document.
-
-        Example:
-            data_dict = {
-                "key1s": [1, 2, 3],
-                "key2s": ["a", "b", "c"],
-                "key3s": None,
-                "key4s": ["x", "y", "z"],
-            }
-
-            results = [
-                {"key1": 1, "key2": "a", "key4": "x"},
-                {"key1": 2, "key2": "b", "key4": "y"},
-                {"key1": 3, "key2": "c", "key4": "z"},
-            ]
-        """
-
-        results = []
-        keys = [key for key in data_dict if data_dict[key] is not None]
-
-        for i in range(len(data_dict[keys[0]])):
-            sub_dict = {}
-            for key in data_dict.keys():
-                if data_dict[key] is not None and len(data_dict[key]) > i:
-                    sub_dict[key[:-1]] = data_dict[key][i]
-            results.append(sub_dict)
-        return results
-
-    def get_docs_by_ids(
-        self, ids: List[ItemID] = None, collection_name: str = None, include=None, **kwargs
-    ) -> List[Document]:
-        """
-        Retrieve documents from the collection of the vector database based on the ids.
-
-        Args:
-            ids: List[ItemID] | A list of document ids. If None, will return all the documents. Default is None.
-            collection_name: str | The name of the collection. Default is None.
-            include: List[str] | The fields to include. Default is None.
-                If None, will include ["metadatas", "documents"], ids will always be included.
-            kwargs: dict | Additional keyword arguments.
-
-        Returns:
-            List[Document] | The results.
-        """
-        collection = self.get_collection(collection_name)
-        include = include if include else ["metadatas", "documents"]
-        results = collection.get(ids, include=include, **kwargs)
-        results = self._chroma_get_results_to_list_documents(results)
-        return results