Merge branch 'main' into improve_update_context

Author: thinkall

OAI_CONFIG_LIST_sample

@@ -13,13 +13,13 @@
         "api_key": "<your Azure OpenAI API key here>",
         "base_url": "<your Azure OpenAI API base here>",
         "api_type": "azure",
-        "api_version": "2024-02-15-preview"
+        "api_version": "2024-02-01"
     },
     {
         "model": "<your Azure OpenAI deployment name>",
         "api_key": "<your Azure OpenAI API key here>",
         "base_url": "<your Azure OpenAI API base here>",
         "api_type": "azure",
-        "api_version": "2024-02-15-preview"
+        "api_version": "2024-02-01"
     }
 ]

autogen/agentchat/contrib/vectordb/pgvectordb.py

@@ -32,10 +32,11 @@ class Collection:
         client: The PGVector client.
         collection_name (str): The name of the collection. Default is "documents".
         embedding_function (Callable): The embedding function used to generate the vector representation.
+            Default is None. SentenceTransformer("all-MiniLM-L6-v2").encode will be used when None.
+            Models can be chosen from:
+            https://huggingface.co/models?library=sentence-transformers
         metadata (Optional[dict]): The metadata of the collection.
         get_or_create (Optional): The flag indicating whether to get or create the collection.
-        model_name: (Optional str) | Sentence embedding model to use. Models can be chosen from:
-            https://huggingface.co/models?library=sentence-transformers
     """
 
     def __init__(
@@ -45,7 +46,6 @@ def __init__(
         embedding_function: Callable = None,
         metadata=None,
         get_or_create=None,
-        model_name="all-MiniLM-L6-v2",
     ):
         """
         Initialize the Collection object.
@@ -56,30 +56,26 @@ def __init__(
             embedding_function: The embedding function used to generate the vector representation.
             metadata: The metadata of the collection.
             get_or_create: The flag indicating whether to get or create the collection.
-            model_name: | Sentence embedding model to use. Models can be chosen from:
-                https://huggingface.co/models?library=sentence-transformers
         Returns:
             None
         """
         self.client = client
-        self.embedding_function = embedding_function
-        self.model_name = model_name
         self.name = self.set_collection_name(collection_name)
         self.require_embeddings_or_documents = False
         self.ids = []
-        try:
-            self.embedding_function = (
-                SentenceTransformer(self.model_name) if embedding_function is None else embedding_function
-            )
-        except Exception as e:
-            logger.error(
-                f"Validate the model name entered: {self.model_name} "
-                f"from https://huggingface.co/models?library=sentence-transformers\nError: {e}"
-            )
-            raise e
+        if embedding_function:
+            self.embedding_function = embedding_function
+        else:
+            self.embedding_function = SentenceTransformer("all-MiniLM-L6-v2").encode
         self.metadata = metadata if metadata else {"hnsw:space": "ip", "hnsw:construction_ef": 32, "hnsw:M": 16}
         self.documents = ""
         self.get_or_create = get_or_create
+        # This will get the model dimension size by computing the embeddings dimensions
+        sentences = [
+            "The weather is lovely today in paradise.",
+        ]
+        embeddings = self.embedding_function(sentences)
+        self.dimension = len(embeddings[0])
 
     def set_collection_name(self, collection_name) -> str:
         name = re.sub("-", "_", collection_name)
@@ -115,14 +111,14 @@ def add(self, ids: List[ItemID], documents: List, embeddings: List = None, metad
         elif metadatas is not None:
             for doc_id, metadata, document in zip(ids, metadatas, documents):
                 metadata = re.sub("'", '"', str(metadata))
-                embedding = self.embedding_function.encode(document)
+                embedding = self.embedding_function(document)
                 sql_values.append((doc_id, metadata, embedding, document))
             sql_string = (
                 f"INSERT INTO {self.name} (id, metadatas, embedding, documents)\n" f"VALUES (%s, %s, %s, %s);\n"
             )
         else:
             for doc_id, document in zip(ids, documents):
-                embedding = self.embedding_function.encode(document)
+                embedding = self.embedding_function(document)
                 sql_values.append((doc_id, document, embedding))
             sql_string = f"INSERT INTO {self.name} (id, documents, embedding)\n" f"VALUES (%s, %s, %s);\n"
         logger.debug(f"Add SQL String:\n{sql_string}\n{sql_values}")
@@ -166,7 +162,7 @@ def upsert(self, ids: List[ItemID], documents: List, embeddings: List = None, me
         elif metadatas is not None:
             for doc_id, metadata, document in zip(ids, metadatas, documents):
                 metadata = re.sub("'", '"', str(metadata))
-                embedding = self.embedding_function.encode(document)
+                embedding = self.embedding_function(document)
                 sql_values.append((doc_id, metadata, embedding, document, metadata, document, embedding))
             sql_string = (
                 f"INSERT INTO {self.name} (id, metadatas, embedding, documents)\n"
@@ -176,7 +172,7 @@ def upsert(self, ids: List[ItemID], documents: List, embeddings: List = None, me
             )
         else:
             for doc_id, document in zip(ids, documents):
-                embedding = self.embedding_function.encode(document)
+                embedding = self.embedding_function(document)
                 sql_values.append((doc_id, document, embedding, document))
             sql_string = (
                 f"INSERT INTO {self.name} (id, documents, embedding)\n"
@@ -304,7 +300,7 @@ def get(
                 )
         except (psycopg.errors.UndefinedTable, psycopg.errors.UndefinedColumn) as e:
             logger.info(f"Error executing select on non-existent table: {self.name}. Creating it instead. Error: {e}")
-            self.create_collection(collection_name=self.name)
+            self.create_collection(collection_name=self.name, dimension=self.dimension)
             logger.info(f"Created table {self.name}")
 
         cursor.close()
@@ -419,7 +415,7 @@ def query(
         cursor = self.client.cursor()
         results = []
         for query_text in query_texts:
-            vector = self.embedding_function.encode(query_text, convert_to_tensor=False).tolist()
+            vector = self.embedding_function(query_text, convert_to_tensor=False).tolist()
             if distance_type.lower() == "cosine":
                 index_function = "<=>"
             elif distance_type.lower() == "euclidean":
@@ -526,22 +522,31 @@ def delete_collection(self, collection_name: Optional[str] = None) -> None:
         cursor.execute(f"DROP TABLE IF EXISTS {self.name}")
         cursor.close()
 
-    def create_collection(self, collection_name: Optional[str] = None) -> None:
+    def create_collection(
+        self, collection_name: Optional[str] = None, dimension: Optional[Union[str, int]] = None
+    ) -> None:
         """
         Create a new collection.
 
         Args:
             collection_name (Optional[str]): The name of the new collection.
+            dimension (Optional[Union[str, int]]): The dimension size of the sentence embedding model
 
         Returns:
             None
         """
         if collection_name:
             self.name = collection_name
+
+        if dimension:
+            self.dimension = dimension
+        elif self.dimension is None:
+            self.dimension = 384
+
         cursor = self.client.cursor()
         cursor.execute(
             f"CREATE TABLE {self.name} ("
-            f"documents text, id CHAR(8) PRIMARY KEY, metadatas JSONB, embedding vector(384));"
+            f"documents text, id CHAR(8) PRIMARY KEY, metadatas JSONB, embedding vector({self.dimension}));"
             f"CREATE INDEX "
             f'ON {self.name} USING hnsw (embedding vector_l2_ops) WITH (m = {self.metadata["hnsw:M"]}, '
             f'ef_construction = {self.metadata["hnsw:construction_ef"]});'
@@ -573,7 +578,6 @@ def __init__(
         connect_timeout: Optional[int] = 10,
         embedding_function: Callable = None,
         metadata: Optional[dict] = None,
-        model_name: Optional[str] = "all-MiniLM-L6-v2",
     ) -> None:
         """
         Initialize the vector database.
@@ -591,15 +595,14 @@ def __init__(
             username: str | The database username to use. Default is None.
             password: str | The database user password to use. Default is None.
             connect_timeout: int | The timeout to set for the connection. Default is 10.
-            embedding_function: Callable | The embedding function used to generate the vector representation
-                of the documents. Default is None.
+            embedding_function: Callable | The embedding function used to generate the vector representation.
+                Default is None. SentenceTransformer("all-MiniLM-L6-v2").encode will be used when None.
+                Models can be chosen from:
+                https://huggingface.co/models?library=sentence-transformers
             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": 16}. Creates Index on table
                 using hnsw (embedding vector_l2_ops) WITH (m = hnsw:M) ef_construction = "hnsw:construction_ef".
                 For more info: https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw
-            model_name: str | Sentence embedding model to use. Models can be chosen from:
-                https://huggingface.co/models?library=sentence-transformers
-
         Returns:
             None
         """
@@ -613,17 +616,10 @@ def __init__(
             password=password,
             connect_timeout=connect_timeout,
         )
-        self.model_name = model_name
-        try:
-            self.embedding_function = (
-                SentenceTransformer(self.model_name) if embedding_function is None else embedding_function
-            )
-        except Exception as e:
-            logger.error(
-                f"Validate the model name entered: {self.model_name} "
-                f"from https://huggingface.co/models?library=sentence-transformers\nError: {e}"
-            )
-            raise e
+        if embedding_function:
+            self.embedding_function = embedding_function
+        else:
+            self.embedding_function = SentenceTransformer("all-MiniLM-L6-v2").encode
         self.metadata = metadata
         register_vector(self.client)
         self.active_collection = None
@@ -738,7 +734,6 @@ def create_collection(
                 embedding_function=self.embedding_function,
                 get_or_create=get_or_create,
                 metadata=self.metadata,
-                model_name=self.model_name,
             )
             collection.set_collection_name(collection_name=collection_name)
             collection.create_collection(collection_name=collection_name)
@@ -751,7 +746,6 @@ def create_collection(
                 embedding_function=self.embedding_function,
                 get_or_create=get_or_create,
                 metadata=self.metadata,
-                model_name=self.model_name,
             )
             collection.set_collection_name(collection_name=collection_name)
             collection.create_collection(collection_name=collection_name)
@@ -765,7 +759,6 @@ def create_collection(
                 embedding_function=self.embedding_function,
                 get_or_create=get_or_create,
                 metadata=self.metadata,
-                model_name=self.model_name,
             )
             collection.set_collection_name(collection_name=collection_name)
             collection.create_collection(collection_name=collection_name)
@@ -797,7 +790,6 @@ def get_collection(self, collection_name: str = None) -> Collection:
                     client=self.client,
                     collection_name=collection_name,
                     embedding_function=self.embedding_function,
-                    model_name=self.model_name,
                 )
         return self.active_collection
 

autogen/agentchat/conversable_agent.py

@@ -31,7 +31,7 @@
 from ..function_utils import get_function_schema, load_basemodels_if_needed, serialize_to_str
 from ..io.base import IOStream
 from ..oai.client import ModelClient, OpenAIWrapper
-from ..runtime_logging import log_event, log_new_agent, logging_enabled
+from ..runtime_logging import log_event, log_function_use, log_new_agent, logging_enabled
 from .agent import Agent, LLMAgent
 from .chat import ChatResult, a_initiate_chats, initiate_chats
 from .utils import consolidate_chat_info, gather_usage_summary
@@ -1357,9 +1357,7 @@ def _generate_oai_reply_from_client(self, llm_client, messages, cache) -> Union[
 
         # TODO: #1143 handle token limit exceeded error
         response = llm_client.create(
-            context=messages[-1].pop("context", None),
-            messages=all_messages,
-            cache=cache,
+            context=messages[-1].pop("context", None), messages=all_messages, cache=cache, agent=self
         )
         extracted_response = llm_client.extract_text_or_completion_object(response)[0]
 
@@ -2528,13 +2526,14 @@ def _wrap_function(self, func: F) -> F:
         @functools.wraps(func)
         def _wrapped_func(*args, **kwargs):
             retval = func(*args, **kwargs)
-
+            log_function_use(self, func, kwargs, retval)
             return serialize_to_str(retval)
 
         @load_basemodels_if_needed
         @functools.wraps(func)
         async def _a_wrapped_func(*args, **kwargs):
             retval = await func(*args, **kwargs)
+            log_function_use(self, func, kwargs, retval)
             return serialize_to_str(retval)
 
         wrapped_func = _a_wrapped_func if inspect.iscoroutinefunction(func) else _wrapped_func

autogen/logger/base_logger.py

@@ -3,14 +3,15 @@
 import sqlite3
 import uuid
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, List, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, TypeVar, Union
 
 from openai import AzureOpenAI, OpenAI
 from openai.types.chat import ChatCompletion
 
 if TYPE_CHECKING:
     from autogen import Agent, ConversableAgent, OpenAIWrapper
 
+F = TypeVar("F", bound=Callable[..., Any])
 ConfigItem = Dict[str, Union[str, List[str]]]
 LLMConfig = Dict[str, Union[None, float, int, ConfigItem, List[ConfigItem]]]
 
@@ -32,6 +33,7 @@ def log_chat_completion(
         invocation_id: uuid.UUID,
         client_id: int,
         wrapper_id: int,
+        source: Union[str, Agent],
         request: Dict[str, Union[float, str, List[Dict[str, str]]]],
         response: Union[str, ChatCompletion],
         is_cached: int,
@@ -49,9 +51,10 @@ def log_chat_completion(
             invocation_id (uuid):               A unique identifier for the invocation to the OpenAIWrapper.create method call
             client_id (int):                    A unique identifier for the underlying OpenAI client instance
             wrapper_id (int):                   A unique identifier for the OpenAIWrapper instance
-            request (dict):                     A dictionary representing the the request or call to the OpenAI client endpoint
+            source (str or Agent):              The source/creator of the event as a string name or an Agent instance
+            request (dict):                     A dictionary representing the request or call to the OpenAI client endpoint
             response (str or ChatCompletion):   The response from OpenAI
-            is_chached (int):                   1 if the response was a cache hit, 0 otherwise
+            is_cached (int):                    1 if the response was a cache hit, 0 otherwise
             cost(float):                        The cost for OpenAI response
             start_time (str):                   A string representing the moment the request was initiated
         """
@@ -104,6 +107,18 @@ def log_new_client(
         """
         ...
 
+    @abstractmethod
+    def log_function_use(self, source: Union[str, Agent], function: F, args: Dict[str, Any], returns: Any) -> None:
+        """
+        Log the use of a registered function (could be a tool)
+
+        Args:
+            source (str or Agent):      The source/creator of the event as a string name or an Agent instance
+            function (F):               The function information
+            args (dict):                The function args to log
+            returns (any):              The return
+        """
+
     @abstractmethod
     def stop(self) -> None:
         """