ag2 0.4.1__py3-none-any.whl → 0.5.0b2__py3-none-any.whl

autogen/io/websockets.py

@@ -1,213 +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 logging
-import ssl
-import threading
-from contextlib import contextmanager
-from functools import partial
-from time import sleep
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Iterator, Optional, Protocol, Union
-
-from .base import IOStream
-
-# Check if the websockets module is available
-try:
-    from websockets.sync.server import serve as ws_serve
-except ImportError as e:
-    _import_error: Optional[ImportError] = e
-else:
-    _import_error = None
-
-
-__all__ = ("IOWebsockets",)
-
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-# The following type and protocols are used to define the ServerConnection and WebSocketServer classes
-# if websockets is not installed, they would be untyped
-Data = Union[str, bytes]
-
-
-class ServerConnection(Protocol):
-    def send(self, message: Union[Data, Iterable[Data]]) -> None:
-        """Send a message to the client.
-
-        Args:
-            message (Union[Data, Iterable[Data]]): The message to send.
-
-        """
-        ...  # pragma: no cover
-
-    def recv(self, timeout: Optional[float] = None) -> Data:
-        """Receive a message from the client.
-
-        Args:
-            timeout (Optional[float], optional): The timeout for the receive operation. Defaults to None.
-
-        Returns:
-            Data: The message received from the client.
-
-        """
-        ...  # pragma: no cover
-
-    def close(self) -> None:
-        """Close the connection."""
-        ...
-
-
-class WebSocketServer(Protocol):
-    def serve_forever(self) -> None:
-        """Run the server forever."""
-        ...  # pragma: no cover
-
-    def shutdown(self) -> None:
-        """Shutdown the server."""
-        ...  # pragma: no cover
-
-    def __enter__(self) -> "WebSocketServer":
-        """Enter the server context."""
-        ...  # pragma: no cover
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        """Exit the server context."""
-        ...  # pragma: no cover
-
-
-class IOWebsockets(IOStream):
-    """A websocket input/output stream."""
-
-    def __init__(self, websocket: ServerConnection) -> None:
-        """Initialize the websocket input/output stream.
-
-        Args:
-            websocket (ServerConnection): The websocket server.
-
-        Raises:
-            ImportError: If the websockets module is not available.
-        """
-        if _import_error is not None:
-            raise _import_error  # pragma: no cover
-
-        self._websocket = websocket
-
-    @staticmethod
-    def _handler(websocket: ServerConnection, on_connect: Callable[["IOWebsockets"], None]) -> None:
-        """The handler function for the websocket server."""
-        logger.info(f" - IOWebsockets._handler(): Client connected on {websocket}")
-        # create a new IOWebsockets instance using the websocket that is create when a client connects
-        try:
-            iowebsocket = IOWebsockets(websocket)
-            with IOStream.set_default(iowebsocket):
-                # call the on_connect function
-                try:
-                    on_connect(iowebsocket)
-                except Exception as e:
-                    logger.warning(f" - IOWebsockets._handler(): Error in on_connect: {e}")
-        except Exception as e:
-            logger.error(f" - IOWebsockets._handler(): Unexpected error in IOWebsockets: {e}")
-
-    @staticmethod
-    @contextmanager
-    def run_server_in_thread(
-        *,
-        host: str = "127.0.0.1",
-        port: int = 8765,
-        on_connect: Callable[["IOWebsockets"], None],
-        ssl_context: Optional[ssl.SSLContext] = None,
-        **kwargs: Any,
-    ) -> Iterator[str]:
-        """Factory function to create a websocket input/output stream.
-
-        Args:
-            host (str, optional): The host to bind the server to. Defaults to "127.0.0.1".
-            port (int, optional): The port to bind the server to. Defaults to 8765.
-            on_connect (Callable[[IOWebsockets], None]): The function to be executed on client connection. Typically creates agents and initiate chat.
-            ssl_context (Optional[ssl.SSLContext], optional): The SSL context to use for secure connections. Defaults to None.
-            kwargs (Any): Additional keyword arguments to pass to the websocket server.
-
-        Yields:
-            str: The URI of the websocket server.
-        """
-        server_dict: Dict[str, WebSocketServer] = {}
-
-        def _run_server() -> None:
-            if _import_error is not None:
-                raise _import_error
-
-            # print(f" - _run_server(): starting server on ws://{host}:{port}", flush=True)
-            with ws_serve(
-                handler=partial(IOWebsockets._handler, on_connect=on_connect),
-                host=host,
-                port=port,
-                ssl_context=ssl_context,
-                **kwargs,
-            ) as server:
-                # print(f" - _run_server(): server {server} started on ws://{host}:{port}", flush=True)
-
-                server_dict["server"] = server
-
-                # runs until the server is shutdown
-                server.serve_forever()
-
-                return
-
-        # start server in a separate thread
-        thread = threading.Thread(target=_run_server)
-        thread.start()
-        try:
-            while "server" not in server_dict:
-                sleep(0.1)
-
-            yield f"ws://{host}:{port}"
-
-        finally:
-            # print(f" - run_server_in_thread(): shutting down server on ws://{host}:{port}", flush=True)
-            # gracefully stop server
-            if "server" in server_dict:
-                # print(f" - run_server_in_thread(): shutting down server {server_dict['server']}", flush=True)
-                server_dict["server"].shutdown()
-
-            # wait for the thread to stop
-            if thread:
-                thread.join()
-
-    @property
-    def websocket(self) -> "ServerConnection":
-        """The URI of the websocket server."""
-        return self._websocket
-
-    def print(self, *objects: Any, sep: str = " ", end: str = "\n", flush: bool = False) -> None:
-        """Print data to the output stream.
-
-        Args:
-            objects (any): The data to print.
-            sep (str, optional): The separator between objects. Defaults to " ".
-            end (str, optional): The end of the output. Defaults to "\n".
-            flush (bool, optional): Whether to flush the output. Defaults to False.
-        """
-        xs = sep.join(map(str, objects)) + end
-        self._websocket.send(xs)
-
-    def input(self, prompt: str = "", *, password: bool = False) -> str:
-        """Read a line from the input stream.
-
-        Args:
-            prompt (str, optional): The prompt to display. Defaults to "".
-            password (bool, optional): Whether to read a password. Defaults to False.
-
-        Returns:
-            str: The line read from the input stream.
-
-        """
-        if prompt != "":
-            self._websocket.send(prompt)
-
-        msg = self._websocket.recv()
-
-        return msg.decode("utf-8") if isinstance(msg, bytes) else msg

autogen/logger/__init__.py

@@ -1,11 +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 .file_logger import FileLogger
-from .logger_factory import LoggerFactory
-from .sqlite_logger import SqliteLogger
-
-__all__ = ("LoggerFactory", "SqliteLogger", "FileLogger")

autogen/logger/base_logger.py

@@ -1,140 +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 __future__ import annotations
-
-import sqlite3
-import uuid
-from abc import ABC, abstractmethod
-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]]]
-
-
-class BaseLogger(ABC):
-    @abstractmethod
-    def start(self) -> str:
-        """
-        Open a connection to the logging database, and start recording.
-
-        Returns:
-            session_id (str):     a unique id for the logging session
-        """
-        ...
-
-    @abstractmethod
-    def log_chat_completion(
-        self,
-        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,
-        cost: float,
-        start_time: str,
-    ) -> None:
-        """
-        Log a chat completion to database.
-
-        In AutoGen, chat completions are somewhat complicated because they are handled by the `autogen.oai.OpenAIWrapper` class.
-        One invocation to `create` can lead to multiple underlying OpenAI calls, depending on the llm_config list used, and
-        any errors or retries.
-
-        Args:
-            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
-            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_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
-        """
-        ...
-
-    @abstractmethod
-    def log_new_agent(self, agent: ConversableAgent, init_args: Dict[str, Any]) -> None:
-        """
-        Log the birth of a new agent.
-
-        Args:
-            agent (ConversableAgent):   The agent to log.
-            init_args (dict):           The arguments passed to the construct the conversable agent
-        """
-        ...
-
-    @abstractmethod
-    def log_event(self, source: Union[str, Agent], name: str, **kwargs: Dict[str, Any]) -> None:
-        """
-        Log an event for an agent.
-
-        Args:
-            source (str or Agent):      The source/creator of the event as a string name or an Agent instance
-            name (str):                 The name of the event
-            kwargs (dict):              The event information to log
-        """
-        ...
-
-    @abstractmethod
-    def log_new_wrapper(self, wrapper: OpenAIWrapper, init_args: Dict[str, Union[LLMConfig, List[LLMConfig]]]) -> None:
-        """
-        Log the birth of a new OpenAIWrapper.
-
-        Args:
-            wrapper (OpenAIWrapper):    The wrapper to log.
-            init_args (dict):           The arguments passed to the construct the wrapper
-        """
-        ...
-
-    @abstractmethod
-    def log_new_client(
-        self, client: Union[AzureOpenAI, OpenAI], wrapper: OpenAIWrapper, init_args: Dict[str, Any]
-    ) -> None:
-        """
-        Log the birth of a new OpenAIWrapper.
-
-        Args:
-            wrapper (OpenAI):           The OpenAI client to log.
-            init_args (dict):           The arguments passed to the construct the 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:
-        """
-        Close the connection to the logging database, and stop logging.
-        """
-        ...
-
-    @abstractmethod
-    def get_connection(self) -> Union[None, sqlite3.Connection]:
-        """
-        Return a connection to the logging database.
-        """
-        ...

autogen/logger/file_logger.py

@@ -1,287 +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 __future__ import annotations
-
-import json
-import logging
-import os
-import threading
-import uuid
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
-
-from openai import AzureOpenAI, OpenAI
-from openai.types.chat import ChatCompletion
-
-from autogen.logger.base_logger import BaseLogger
-from autogen.logger.logger_utils import get_current_ts, to_dict
-
-from .base_logger import LLMConfig
-
-if TYPE_CHECKING:
-    from autogen import Agent, ConversableAgent, OpenAIWrapper
-    from autogen.oai.anthropic import AnthropicClient
-    from autogen.oai.bedrock import BedrockClient
-    from autogen.oai.cerebras import CerebrasClient
-    from autogen.oai.cohere import CohereClient
-    from autogen.oai.gemini import GeminiClient
-    from autogen.oai.groq import GroqClient
-    from autogen.oai.mistral import MistralAIClient
-    from autogen.oai.ollama import OllamaClient
-    from autogen.oai.together import TogetherClient
-
-logger = logging.getLogger(__name__)
-
-F = TypeVar("F", bound=Callable[..., Any])
-
-__all__ = ("FileLogger",)
-
-
-def safe_serialize(obj: Any) -> str:
-    def default(o: Any) -> str:
-        if hasattr(o, "to_json"):
-            return str(o.to_json())
-        else:
-            return f"<<non-serializable: {type(o).__qualname__}>>"
-
-    return json.dumps(obj, default=default)
-
-
-class FileLogger(BaseLogger):
-    def __init__(self, config: Dict[str, Any]):
-        self.config = config
-        self.session_id = str(uuid.uuid4())
-
-        curr_dir = os.getcwd()
-        self.log_dir = os.path.join(curr_dir, "autogen_logs")
-        os.makedirs(self.log_dir, exist_ok=True)
-
-        self.log_file = os.path.join(self.log_dir, self.config.get("filename", "runtime.log"))
-        try:
-            with open(self.log_file, "a"):
-                pass
-        except Exception as e:
-            logger.error(f"[file_logger] Failed to create logging file: {e}")
-
-        self.logger = logging.getLogger(__name__)
-        self.logger.setLevel(logging.INFO)
-        file_handler = logging.FileHandler(self.log_file)
-        self.logger.addHandler(file_handler)
-
-    def start(self) -> str:
-        """Start the logger and return the session_id."""
-        try:
-            self.logger.info(f"Started new session with Session ID: {self.session_id}")
-        except Exception as e:
-            logger.error(f"[file_logger] Failed to create logging file: {e}")
-        finally:
-            return self.session_id
-
-    def log_chat_completion(
-        self,
-        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,
-        cost: float,
-        start_time: str,
-    ) -> None:
-        """
-        Log a chat completion.
-        """
-        thread_id = threading.get_ident()
-        source_name = None
-        if isinstance(source, str):
-            source_name = source
-        else:
-            source_name = source.name
-        try:
-            log_data = json.dumps(
-                {
-                    "invocation_id": str(invocation_id),
-                    "client_id": client_id,
-                    "wrapper_id": wrapper_id,
-                    "request": to_dict(request),
-                    "response": str(response),
-                    "is_cached": is_cached,
-                    "cost": cost,
-                    "start_time": start_time,
-                    "end_time": get_current_ts(),
-                    "thread_id": thread_id,
-                    "source_name": source_name,
-                }
-            )
-
-            self.logger.info(log_data)
-        except Exception as e:
-            self.logger.error(f"[file_logger] Failed to log chat completion: {e}")
-
-    def log_new_agent(self, agent: ConversableAgent, init_args: Dict[str, Any] = {}) -> None:
-        """
-        Log a new agent instance.
-        """
-        thread_id = threading.get_ident()
-
-        try:
-            log_data = json.dumps(
-                {
-                    "id": id(agent),
-                    "agent_name": agent.name if hasattr(agent, "name") and agent.name is not None else "",
-                    "wrapper_id": to_dict(
-                        agent.client.wrapper_id if hasattr(agent, "client") and agent.client is not None else ""
-                    ),
-                    "session_id": self.session_id,
-                    "current_time": get_current_ts(),
-                    "agent_type": type(agent).__name__,
-                    "args": to_dict(init_args),
-                    "thread_id": thread_id,
-                }
-            )
-            self.logger.info(log_data)
-        except Exception as e:
-            self.logger.error(f"[file_logger] Failed to log new agent: {e}")
-
-    def log_event(self, source: Union[str, Agent], name: str, **kwargs: Dict[str, Any]) -> None:
-        """
-        Log an event from an agent or a string source.
-        """
-        from autogen import Agent
-
-        # This takes an object o as input and returns a string. If the object o cannot be serialized, instead of raising an error,
-        # it returns a string indicating that the object is non-serializable, along with its type's qualified name obtained using __qualname__.
-        json_args = json.dumps(kwargs, default=lambda o: f"<<non-serializable: {type(o).__qualname__}>>")
-        thread_id = threading.get_ident()
-
-        if isinstance(source, Agent):
-            try:
-                log_data = json.dumps(
-                    {
-                        "source_id": id(source),
-                        "source_name": str(source.name) if hasattr(source, "name") else source,
-                        "event_name": name,
-                        "agent_module": source.__module__,
-                        "agent_class": source.__class__.__name__,
-                        "json_state": json_args,
-                        "timestamp": get_current_ts(),
-                        "thread_id": thread_id,
-                    }
-                )
-                self.logger.info(log_data)
-            except Exception as e:
-                self.logger.error(f"[file_logger] Failed to log event {e}")
-        else:
-            try:
-                log_data = json.dumps(
-                    {
-                        "source_id": id(source),
-                        "source_name": str(source.name) if hasattr(source, "name") else source,
-                        "event_name": name,
-                        "json_state": json_args,
-                        "timestamp": get_current_ts(),
-                        "thread_id": thread_id,
-                    }
-                )
-                self.logger.info(log_data)
-            except Exception as e:
-                self.logger.error(f"[file_logger] Failed to log event {e}")
-
-    def log_new_wrapper(
-        self, wrapper: OpenAIWrapper, init_args: Dict[str, Union[LLMConfig, List[LLMConfig]]] = {}
-    ) -> None:
-        """
-        Log a new wrapper instance.
-        """
-        thread_id = threading.get_ident()
-
-        try:
-            log_data = json.dumps(
-                {
-                    "wrapper_id": id(wrapper),
-                    "session_id": self.session_id,
-                    "json_state": json.dumps(init_args),
-                    "timestamp": get_current_ts(),
-                    "thread_id": thread_id,
-                }
-            )
-            self.logger.info(log_data)
-        except Exception as e:
-            self.logger.error(f"[file_logger] Failed to log event {e}")
-
-    def log_new_client(
-        self,
-        client: (
-            AzureOpenAI
-            | OpenAI
-            | CerebrasClient
-            | GeminiClient
-            | AnthropicClient
-            | MistralAIClient
-            | TogetherClient
-            | GroqClient
-            | CohereClient
-            | OllamaClient
-            | BedrockClient
-        ),
-        wrapper: OpenAIWrapper,
-        init_args: Dict[str, Any],
-    ) -> None:
-        """
-        Log a new client instance.
-        """
-        thread_id = threading.get_ident()
-
-        try:
-            log_data = json.dumps(
-                {
-                    "client_id": id(client),
-                    "wrapper_id": id(wrapper),
-                    "session_id": self.session_id,
-                    "class": type(client).__name__,
-                    "json_state": json.dumps(init_args),
-                    "timestamp": get_current_ts(),
-                    "thread_id": thread_id,
-                }
-            )
-            self.logger.info(log_data)
-        except Exception as e:
-            self.logger.error(f"[file_logger] Failed to log event {e}")
-
-    def log_function_use(self, source: Union[str, Agent], function: F, args: Dict[str, Any], returns: Any) -> None:
-        """
-        Log a registered function(can be a tool) use from an agent or a string source.
-        """
-        thread_id = threading.get_ident()
-
-        try:
-            log_data = json.dumps(
-                {
-                    "source_id": id(source),
-                    "source_name": str(source.name) if hasattr(source, "name") else source,
-                    "agent_module": source.__module__,
-                    "agent_class": source.__class__.__name__,
-                    "timestamp": get_current_ts(),
-                    "thread_id": thread_id,
-                    "input_args": safe_serialize(args),
-                    "returns": safe_serialize(returns),
-                }
-            )
-            self.logger.info(log_data)
-        except Exception as e:
-            self.logger.error(f"[file_logger] Failed to log event {e}")
-
-    def get_connection(self) -> None:
-        """Method is intentionally left blank because there is no specific connection needed for the FileLogger."""
-        pass
-
-    def stop(self) -> None:
-        """Close the file handler and remove it from the logger."""
-        for handler in self.logger.handlers:
-            if isinstance(handler, logging.FileHandler):
-                handler.close()
-                self.logger.removeHandler(handler)

autogen/logger/logger_factory.py

@@ -1,29 +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, Dict, Literal, Optional
-
-from autogen.logger.base_logger import BaseLogger
-from autogen.logger.file_logger import FileLogger
-from autogen.logger.sqlite_logger import SqliteLogger
-
-__all__ = ("LoggerFactory",)
-
-
-class LoggerFactory:
-    @staticmethod
-    def get_logger(
-        logger_type: Literal["sqlite", "file"] = "sqlite", config: Optional[Dict[str, Any]] = None
-    ) -> BaseLogger:
-        if config is None:
-            config = {}
-
-        if logger_type == "sqlite":
-            return SqliteLogger(config)
-        elif logger_type == "file":
-            return FileLogger(config)
-        else:
-            raise ValueError(f"[logger_factory] Unknown logger type: {logger_type}")