ag2 0.9.7__py3-none-any.whl → 0.9.9__py3-none-any.whl

autogen/io/base.py

@@ -8,7 +8,7 @@ import logging
 from collections.abc import Iterator
 from contextlib import contextmanager
 from contextvars import ContextVar
-from typing import Any, Optional, Protocol, Union, runtime_checkable
+from typing import Any, Protocol, runtime_checkable
 
 from ..doc_utils import export_module
 from ..events.base_event import BaseEvent
@@ -87,7 +87,7 @@ class AsyncIOStreamProtocol(AsyncInputStream, OutputStream, Protocol):
     """A protocol for input/output streams."""
 
 
-iostream_union = Union[IOStreamProtocol, AsyncIOStreamProtocol]
+iostream_union = IOStreamProtocol | AsyncIOStreamProtocol
 
 
 @export_module("autogen.io")
@@ -95,9 +95,9 @@ class IOStream:
     """A protocol for input/output streams."""
 
     # ContextVar must be used in multithreaded or async environments
-    _default_io_stream: ContextVar[Optional[iostream_union]] = ContextVar("default_iostream", default=None)
+    _default_io_stream: ContextVar[iostream_union | None] = ContextVar("default_iostream", default=None)
     _default_io_stream.set(None)
-    _global_default: Optional[iostream_union] = None
+    _global_default: iostream_union | None = None
 
     @staticmethod
     def set_global_default(stream: iostream_union) -> None:
@@ -135,7 +135,7 @@ class IOStream:
 
     @staticmethod
     @contextmanager
-    def set_default(stream: Optional[iostream_union]) -> Iterator[None]:
+    def set_default(stream: iostream_union | None) -> Iterator[None]:
         """Set the default input/output stream.
 
         Args:

autogen/io/run_response.py

@@ -7,7 +7,8 @@
 
 import queue
 from asyncio import Queue as AsyncQueue
-from typing import Any, AsyncIterable, Dict, Iterable, Optional, Protocol, Sequence, Union
+from collections.abc import AsyncIterable, Iterable, Sequence
+from typing import Any, Optional, Protocol
 from uuid import UUID, uuid4
 
 from pydantic import BaseModel, Field
@@ -46,7 +47,7 @@ class Usage(BaseModel):
 
 class CostBreakdown(BaseModel):
     total_cost: float
-    models: Dict[str, Usage] = Field(default_factory=dict)
+    models: dict[str, Usage] = Field(default_factory=dict)
 
     @classmethod
     def from_raw(cls, data: dict[str, Any]) -> "CostBreakdown":
@@ -79,18 +80,18 @@ class RunResponseProtocol(RunInfoProtocol, Protocol):
     def messages(self) -> Iterable[Message]: ...
 
     @property
-    def summary(self) -> Optional[str]: ...
+    def summary(self) -> str | None: ...
 
     @property
-    def context_variables(self) -> Optional[ContextVariables]: ...
+    def context_variables(self) -> ContextVariables | None: ...
 
     @property
-    def last_speaker(self) -> Optional[str]: ...
+    def last_speaker(self) -> str | None: ...
 
     @property
-    def cost(self) -> Optional[Cost]: ...
+    def cost(self) -> Cost | None: ...
 
-    def process(self, processor: Optional[EventProcessorProtocol] = None) -> None: ...
+    def process(self, processor: EventProcessorProtocol | None = None) -> None: ...
 
     def set_ui_tools(self, tools: list[Tool]) -> None: ...
 
@@ -103,18 +104,18 @@ class AsyncRunResponseProtocol(RunInfoProtocol, Protocol):
     async def messages(self) -> Iterable[Message]: ...
 
     @property
-    async def summary(self) -> Optional[str]: ...
+    async def summary(self) -> str | None: ...
 
     @property
-    async def context_variables(self) -> Optional[ContextVariables]: ...
+    async def context_variables(self) -> ContextVariables | None: ...
 
     @property
-    async def last_speaker(self) -> Optional[str]: ...
+    async def last_speaker(self) -> str | None: ...
 
     @property
-    async def cost(self) -> Optional[Cost]: ...
+    async def cost(self) -> Cost | None: ...
 
-    async def process(self, processor: Optional[AsyncEventProcessorProtocol] = None) -> None: ...
+    async def process(self, processor: AsyncEventProcessorProtocol | None = None) -> None: ...
 
     def set_ui_tools(self, tools: list[Tool]) -> None: ...
 
@@ -123,12 +124,12 @@ class RunResponse:
     def __init__(self, iostream: ThreadIOStream, agents: list[Agent]):
         self.iostream = iostream
         self.agents = agents
-        self._summary: Optional[str] = None
+        self._summary: str | None = None
         self._messages: Sequence[LLMMessageType] = []
         self._uuid = uuid4()
-        self._context_variables: Optional[ContextVariables] = None
-        self._last_speaker: Optional[str] = None
-        self._cost: Optional[Cost] = None
+        self._context_variables: ContextVariables | None = None
+        self._last_speaker: str | None = None
+        self._cost: Cost | None = None
 
     def _queue_generator(self, q: queue.Queue) -> Iterable[BaseEvent]:  # type: ignore[type-arg]
         """A generator to yield items from the queue until the termination message is found."""
@@ -164,7 +165,7 @@ class RunResponse:
         return self._messages
 
     @property
-    def summary(self) -> Optional[str]:
+    def summary(self) -> str | None:
         return self._summary
 
     @property
@@ -176,25 +177,25 @@ class RunResponse:
         return self._uuid
 
     @property
-    def context_variables(self) -> Optional[ContextVariables]:
+    def context_variables(self) -> ContextVariables | None:
         return self._context_variables
 
     @property
-    def last_speaker(self) -> Optional[str]:
+    def last_speaker(self) -> str | None:
         return self._last_speaker
 
     @property
-    def cost(self) -> Optional[Cost]:
+    def cost(self) -> Cost | None:
         return self._cost
 
     @cost.setter
-    def cost(self, value: Union[Cost, dict[str, Any]]) -> None:
+    def cost(self, value: Cost | dict[str, Any]) -> None:
         if isinstance(value, dict):
             self._cost = Cost.from_raw(value)
         else:
             self._cost = value
 
-    def process(self, processor: Optional[EventProcessorProtocol] = None) -> None:
+    def process(self, processor: EventProcessorProtocol | None = None) -> None:
         processor = processor or ConsoleEventProcessor()
         processor.process(self)
 
@@ -208,12 +209,12 @@ class AsyncRunResponse:
     def __init__(self, iostream: AsyncThreadIOStream, agents: list[Agent]):
         self.iostream = iostream
         self.agents = agents
-        self._summary: Optional[str] = None
+        self._summary: str | None = None
         self._messages: Sequence[LLMMessageType] = []
         self._uuid = uuid4()
-        self._context_variables: Optional[ContextVariables] = None
-        self._last_speaker: Optional[str] = None
-        self._cost: Optional[Cost] = None
+        self._context_variables: ContextVariables | None = None
+        self._last_speaker: str | None = None
+        self._cost: Cost | None = None
 
     async def _queue_generator(self, q: AsyncQueue[Any]) -> AsyncIterable[BaseEvent]:  # type: ignore[type-arg]
         """A generator to yield items from the queue until the termination message is found."""
@@ -253,7 +254,7 @@ class AsyncRunResponse:
         return self._messages
 
     @property
-    async def summary(self) -> Optional[str]:
+    async def summary(self) -> str | None:
         return self._summary
 
     @property
@@ -265,25 +266,25 @@ class AsyncRunResponse:
         return self._uuid
 
     @property
-    async def context_variables(self) -> Optional[ContextVariables]:
+    async def context_variables(self) -> ContextVariables | None:
         return self._context_variables
 
     @property
-    async def last_speaker(self) -> Optional[str]:
+    async def last_speaker(self) -> str | None:
         return self._last_speaker
 
     @property
-    async def cost(self) -> Optional[Cost]:
+    async def cost(self) -> Cost | None:
         return self._cost
 
     @cost.setter
-    def cost(self, value: Union[Cost, dict[str, Any]]) -> None:
+    def cost(self, value: Cost | dict[str, Any]) -> None:
         if isinstance(value, dict):
             self._cost = Cost.from_raw(value)
         else:
             self._cost = value
 
-    async def process(self, processor: Optional[AsyncEventProcessorProtocol] = None) -> None:
+    async def process(self, processor: AsyncEventProcessorProtocol | None = None) -> None:
         processor = processor or AsyncConsoleEventProcessor()
         await processor.process(self)
 

autogen/io/websockets.py

@@ -7,10 +7,11 @@
 import logging
 import ssl
 import threading
+from collections.abc import Callable, Iterable, Iterator
 from contextlib import contextmanager
 from functools import partial
 from time import sleep
-from typing import Any, Callable, Iterable, Iterator, Optional, Protocol, Union
+from typing import Any, Protocol
 
 from ..doc_utils import export_module
 from ..events.base_event import BaseEvent
@@ -30,11 +31,11 @@ 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]
+Data = str | bytes
 
 
 class ServerConnection(Protocol):
-    def send(self, message: Union[Data, Iterable[Data]]) -> None:
+    def send(self, message: Data | Iterable[Data]) -> None:
         """Send a message to the client.
 
         Args:
@@ -43,7 +44,7 @@ class ServerConnection(Protocol):
         """
         ...  # pragma: no cover
 
-    def recv(self, timeout: Optional[float] = None) -> Data:
+    def recv(self, timeout: float | None = None) -> Data:
         """Receive a message from the client.
 
         Args:
@@ -114,7 +115,7 @@ class IOWebsockets(IOStream):
         host: str = "127.0.0.1",
         port: int = 8765,
         on_connect: Callable[["IOWebsockets"], None],
-        ssl_context: Optional[ssl.SSLContext] = None,
+        ssl_context: ssl.SSLContext | None = None,
         **kwargs: Any,
     ) -> Iterator[str]:
         """Factory function to create a websocket input/output stream.

autogen/json_utils.py

@@ -14,8 +14,7 @@ __all__ = ["resolve_json_references"]
 
 @require_optional_import("jsonschema", "gemini")
 def resolve_json_references(schema: dict[str, Any]) -> dict[str, Any]:
-    """
-    Resolve JSON references in the given schema.
+    """Resolve JSON references in the given schema.
 
     Args:
         schema (dict): The JSON schema with references.

autogen/llm_config/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .client import ModelClient
+from .config import LLMConfig
+
+__all__ = (
+    "LLMConfig",
+    "ModelClient",
+)

autogen/llm_config/client.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# 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, Protocol
+
+from ..doc_utils import export_module
+
+
+@export_module("autogen")
+class ModelClient(Protocol):
+    """A client class must implement the following methods:
+    - create must return a response object that implements the ModelClientResponseProtocol
+    - cost must return the cost of the response
+    - get_usage must return a dict with the following keys:
+        - prompt_tokens
+        - completion_tokens
+        - total_tokens
+        - cost
+        - model
+
+    This class is used to create a client that can be used by OpenAIWrapper.
+    The response returned from create must adhere to the ModelClientResponseProtocol but can be extended however needed.
+    The message_retrieval method must be implemented to return a list of str or a list of messages from the response.
+    """
+
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
+    class ModelClientResponseProtocol(Protocol):
+        class Choice(Protocol):
+            class Message(Protocol):
+                content: str | dict[str, Any]
+
+            message: Message
+
+        choices: list[Choice]
+        model: str
+
+    def create(self, params: dict[str, Any]) -> ModelClientResponseProtocol: ...  # pragma: no cover
+
+    def message_retrieval(
+        self, response: ModelClientResponseProtocol
+    ) -> list[str] | list["ModelClient.ModelClientResponseProtocol.Choice.Message"]:
+        """Retrieve and return a list of strings or a list of Choice.Message from the response.
+
+        NOTE: if a list of Choice.Message is returned, it currently needs to contain the fields of OpenAI's ChatCompletion Message object,
+        since that is expected for function or tool calling in the rest of the codebase at the moment, unless a custom agent is being used.
+        """
+        ...  # pragma: no cover
+
+    def cost(self, response: ModelClientResponseProtocol) -> float: ...  # pragma: no cover
+
+    @staticmethod
+    def get_usage(response: ModelClientResponseProtocol) -> dict[str, Any]:
+        """Return usage summary of the response using RESPONSE_USAGE_KEYS."""
+        ...  # pragma: no cover