ag2 0.3.2b2__py3-none-any.whl

autogen/agentchat/contrib/graph_rag/graph_query_engine.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass, field
+from typing import List, Optional, Protocol
+
+from .document import Document
+
+
+@dataclass
+class GraphStoreQueryResult:
+    """
+    A wrapper of graph store query results.
+
+    answer: human readable answer to question/query.
+    results: intermediate results to question/query, e.g. node entities.
+    """
+
+    answer: Optional[str] = None
+    results: list = field(default_factory=list)
+
+
+class GraphQueryEngine(Protocol):
+    """An abstract base class that represents a graph query engine on top of a underlying graph database.
+
+    This interface defines the basic methods for graph rag.
+    """
+
+    def init_db(self, input_doc: List[Document] | None = None):
+        """
+        This method initializes graph database with the input documents or records.
+        Usually, it takes the following steps,
+        1. connecting to a graph database.
+        2. extract graph nodes, edges based on input data, graph schema and etc.
+        3. build indexes etc.
+
+        Args:
+        input_doc: a list of input documents that are used to build the graph in database.
+
+        """
+        pass
+
+    def add_records(self, new_records: List) -> bool:
+        """
+        Add new records to the underlying database and add to the graph if required.
+        """
+        pass
+
+    def query(self, question: str, n_results: int = 1, **kwargs) -> GraphStoreQueryResult:
+        """
+        This method transform a string format question into database query and return the result.
+        """
+        pass

autogen/agentchat/contrib/graph_rag/graph_rag_capability.py

@@ -0,0 +1,56 @@
+from autogen.agentchat.contrib.capabilities.agent_capability import AgentCapability
+from autogen.agentchat.conversable_agent import ConversableAgent
+
+from .graph_query_engine import GraphQueryEngine
+
+
+class GraphRagCapability(AgentCapability):
+    """
+    A graph rag capability uses a graph query engine to give a conversable agent the graph rag ability.
+
+    An agent class with graph rag capability could
+    1. create a graph in the underlying database with input documents.
+    2. retrieved relevant information based on messages received by the agent.
+    3. generate answers from retrieved information and send messages back.
+
+    For example,
+    graph_query_engine = GraphQueryEngine(...)
+    graph_query_engine.init_db([Document(doc1), Document(doc2), ...])
+
+    graph_rag_agent = ConversableAgent(
+        name="graph_rag_agent",
+        max_consecutive_auto_reply=3,
+        ...
+    )
+    graph_rag_capability = GraphRagCapbility(graph_query_engine)
+    graph_rag_capability.add_to_agent(graph_rag_agent)
+
+    user_proxy = UserProxyAgent(
+        name="user_proxy",
+        code_execution_config=False,
+        is_termination_msg=lambda msg: "TERMINATE" in msg["content"],
+        human_input_mode="ALWAYS",
+    )
+    user_proxy.initiate_chat(graph_rag_agent, message="Name a few actors who've played in 'The Matrix'")
+
+    # ChatResult(
+        # chat_id=None,
+        # chat_history=[
+            # {'content': 'Name a few actors who've played in \'The Matrix\'', 'role': 'graph_rag_agent'},
+            # {'content': 'A few actors who have played in The Matrix are:
+            #   - Keanu Reeves
+            #   - Laurence Fishburne
+            #   - Carrie-Anne Moss
+            #   - Hugo Weaving',
+            #   'role': 'user_proxy'},
+        # ...)
+
+    """
+
+    def __init__(self, query_engine: GraphQueryEngine):
+        """
+        initialize graph rag capability with a graph query engine
+        """
+        ...
+
+    def add_to_agent(self, agent: ConversableAgent): ...

autogen/agentchat/contrib/img_utils.py

@@ -0,0 +1,390 @@
+# 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 base64
+import copy
+import os
+import re
+from io import BytesIO
+from math import ceil
+from typing import Dict, List, Tuple, Union
+
+import requests
+from PIL import Image
+
+from autogen.agentchat import utils
+
+# Parameters for token counting for images for different models
+MODEL_PARAMS = {
+    "gpt-4-vision": {
+        "max_edge": 2048,
+        "min_edge": 768,
+        "tile_size": 512,
+        "base_token_count": 85,
+        "token_multiplier": 170,
+    },
+    "gpt-4o-mini": {
+        "max_edge": 2048,
+        "min_edge": 768,
+        "tile_size": 512,
+        "base_token_count": 2833,
+        "token_multiplier": 5667,
+    },
+    "gpt-4o": {"max_edge": 2048, "min_edge": 768, "tile_size": 512, "base_token_count": 85, "token_multiplier": 170},
+}
+
+
+def get_pil_image(image_file: Union[str, Image.Image]) -> Image.Image:
+    """
+    Loads an image from a file and returns a PIL Image object.
+
+    Parameters:
+        image_file (str, or Image): The filename, URL, URI, or base64 string of the image file.
+
+    Returns:
+        Image.Image: The PIL Image object.
+    """
+    if isinstance(image_file, Image.Image):
+        # Already a PIL Image object
+        return image_file
+
+    # Remove quotes if existed
+    if image_file.startswith('"') and image_file.endswith('"'):
+        image_file = image_file[1:-1]
+    if image_file.startswith("'") and image_file.endswith("'"):
+        image_file = image_file[1:-1]
+
+    if image_file.startswith("http://") or image_file.startswith("https://"):
+        # A URL file
+        response = requests.get(image_file)
+        content = BytesIO(response.content)
+        image = Image.open(content)
+    elif re.match(r"data:image/(?:png|jpeg);base64,", image_file):
+        # A URI. Remove the prefix and decode the base64 string.
+        base64_data = re.sub(r"data:image/(?:png|jpeg);base64,", "", image_file)
+        image = _to_pil(base64_data)
+    elif os.path.exists(image_file):
+        # A local file
+        image = Image.open(image_file)
+    else:
+        # base64 encoded string
+        image = _to_pil(image_file)
+
+    return image.convert("RGB")
+
+
+def get_image_data(image_file: Union[str, Image.Image], use_b64=True) -> bytes:
+    """
+    Loads an image and returns its data either as raw bytes or in base64-encoded format.
+
+    This function first loads an image from the specified file, URL, or base64 string using
+    the `get_pil_image` function. It then saves this image in memory in PNG format and
+    retrieves its binary content. Depending on the `use_b64` flag, this binary content is
+    either returned directly or as a base64-encoded string.
+
+    Parameters:
+        image_file (str, or Image): The path to the image file, a URL to an image, or a base64-encoded
+                          string of the image.
+        use_b64 (bool): If True, the function returns a base64-encoded string of the image data.
+                        If False, it returns the raw byte data of the image. Defaults to True.
+
+    Returns:
+        bytes: The image data in raw bytes if `use_b64` is False, or a base64-encoded string
+               if `use_b64` is True.
+    """
+    image = get_pil_image(image_file)
+
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    content = buffered.getvalue()
+
+    if use_b64:
+        return base64.b64encode(content).decode("utf-8")
+    else:
+        return content
+
+
+def llava_formatter(prompt: str, order_image_tokens: bool = False) -> Tuple[str, List[str]]:
+    """
+    Formats the input prompt by replacing image tags and returns the new prompt along with image locations.
+
+    Parameters:
+        - prompt (str): The input string that may contain image tags like <img ...>.
+        - order_image_tokens (bool, optional): Whether to order the image tokens with numbers.
+            It will be useful for GPT-4V. Defaults to False.
+
+    Returns:
+        - Tuple[str, List[str]]: A tuple containing the formatted string and a list of images (loaded in b64 format).
+    """
+
+    # Initialize variables
+    new_prompt = prompt
+    image_locations = []
+    images = []
+    image_count = 0
+
+    # Regular expression pattern for matching <img ...> tags
+    img_tag_pattern = re.compile(r"<img ([^>]+)>")
+
+    # Find all image tags
+    for match in img_tag_pattern.finditer(prompt):
+        image_location = match.group(1)
+
+        try:
+            img_data = get_image_data(image_location)
+        except Exception as e:
+            # Remove the token
+            print(f"Warning! Unable to load image from {image_location}, because of {e}")
+            new_prompt = new_prompt.replace(match.group(0), "", 1)
+            continue
+
+        image_locations.append(image_location)
+        images.append(img_data)
+
+        # Increment the image count and replace the tag in the prompt
+        new_token = f"<image {image_count}>" if order_image_tokens else "<image>"
+
+        new_prompt = new_prompt.replace(match.group(0), new_token, 1)
+        image_count += 1
+
+    return new_prompt, images
+
+
+def pil_to_data_uri(image: Image.Image) -> str:
+    """
+    Converts a PIL Image object to a data URI.
+
+    Parameters:
+        image (Image.Image): The PIL Image object.
+
+    Returns:
+        str: The data URI string.
+    """
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    content = buffered.getvalue()
+    return convert_base64_to_data_uri(base64.b64encode(content).decode("utf-8"))
+
+
+def convert_base64_to_data_uri(base64_image):
+    def _get_mime_type_from_data_uri(base64_image):
+        # Decode the base64 string
+        image_data = base64.b64decode(base64_image)
+        # Check the first few bytes for known signatures
+        if image_data.startswith(b"\xff\xd8\xff"):
+            return "image/jpeg"
+        elif image_data.startswith(b"\x89PNG\r\n\x1a\n"):
+            return "image/png"
+        elif image_data.startswith(b"GIF87a") or image_data.startswith(b"GIF89a"):
+            return "image/gif"
+        elif image_data.startswith(b"RIFF") and image_data[8:12] == b"WEBP":
+            return "image/webp"
+        return "image/jpeg"  # use jpeg for unknown formats, best guess.
+
+    mime_type = _get_mime_type_from_data_uri(base64_image)
+    data_uri = f"data:{mime_type};base64,{base64_image}"
+    return data_uri
+
+
+def gpt4v_formatter(prompt: str, img_format: str = "uri") -> List[Union[str, dict]]:
+    """
+    Formats the input prompt by replacing image tags and returns a list of text and images.
+
+    Args:
+        - prompt (str): The input string that may contain image tags like <img ...>.
+        - img_format (str): what image format should be used. One of "uri", "url", "pil".
+
+    Returns:
+        - List[Union[str, dict]]: A list of alternating text and image dictionary items.
+    """
+    assert img_format in ["uri", "url", "pil"]
+
+    output = []
+    last_index = 0
+    image_count = 0
+
+    # Find all image tags
+    for parsed_tag in utils.parse_tags_from_content("img", prompt):
+        image_location = parsed_tag["attr"]["src"]
+        try:
+            if img_format == "pil":
+                img_data = get_pil_image(image_location)
+            elif img_format == "uri":
+                img_data = get_image_data(image_location)
+                img_data = convert_base64_to_data_uri(img_data)
+            elif img_format == "url":
+                img_data = image_location
+            else:
+                raise ValueError(f"Unknown image format {img_format}")
+        except Exception as e:
+            # Warning and skip this token
+            print(f"Warning! Unable to load image from {image_location}, because {e}")
+            continue
+
+        # Add text before this image tag to output list
+        output.append({"type": "text", "text": prompt[last_index : parsed_tag["match"].start()]})
+
+        # Add image data to output list
+        output.append({"type": "image_url", "image_url": {"url": img_data}})
+
+        last_index = parsed_tag["match"].end()
+        image_count += 1
+
+    # Add remaining text to output list
+    output.append({"type": "text", "text": prompt[last_index:]})
+    return output
+
+
+def extract_img_paths(paragraph: str) -> list:
+    """
+    Extract image paths (URLs or local paths) from a text paragraph.
+
+    Parameters:
+        paragraph (str): The input text paragraph.
+
+    Returns:
+        list: A list of extracted image paths.
+    """
+    # Regular expression to match image URLs and file paths
+    img_path_pattern = re.compile(
+        r"\b(?:http[s]?://\S+\.(?:jpg|jpeg|png|gif|bmp)|\S+\.(?:jpg|jpeg|png|gif|bmp))\b", re.IGNORECASE
+    )
+
+    # Find all matches in the paragraph
+    img_paths = re.findall(img_path_pattern, paragraph)
+    return img_paths
+
+
+def _to_pil(data: str) -> Image.Image:
+    """
+    Converts a base64 encoded image data string to a PIL Image object.
+
+    This function first decodes the base64 encoded string to bytes, then creates a BytesIO object from the bytes,
+    and finally creates and returns a PIL Image object from the BytesIO object.
+
+    Parameters:
+        data (str): The encoded image data string.
+
+    Returns:
+        Image.Image: The PIL Image object created from the input data.
+    """
+    return Image.open(BytesIO(base64.b64decode(data)))
+
+
+def message_formatter_pil_to_b64(messages: List[Dict]) -> List[Dict]:
+    """
+    Converts the PIL image URLs in the messages to base64 encoded data URIs.
+
+    This function iterates over a list of message dictionaries. For each message,
+    if it contains a 'content' key with a list of items, it looks for items
+    with an 'image_url' key. The function then converts the PIL image URL
+    (pointed to by 'image_url') to a base64 encoded data URI.
+
+    Parameters:
+        messages (List[Dict]): A list of message dictionaries. Each dictionary
+                               may contain a 'content' key with a list of items,
+                               some of which might be image URLs.
+
+    Returns:
+        List[Dict]: A new list of message dictionaries with PIL image URLs in the
+                    'image_url' key converted to base64 encoded data URIs.
+
+    Example Input:
+        [
+            {'content': [{'type': 'text', 'text': 'You are a helpful AI assistant.'}], 'role': 'system'},
+            {'content': [
+                {'type': 'text', 'text': "What's the breed of this dog here? \n"},
+                {'type': 'image_url', 'image_url': {'url': a PIL.Image.Image}},
+                {'type': 'text', 'text': '.'}],
+            'role': 'user'}
+        ]
+
+    Example Output:
+        [
+            {'content': [{'type': 'text', 'text': 'You are a helpful AI assistant.'}], 'role': 'system'},
+            {'content': [
+                {'type': 'text', 'text': "What's the breed of this dog here? \n"},
+                {'type': 'image_url', 'image_url': {'url': a B64 Image}},
+                {'type': 'text', 'text': '.'}],
+            'role': 'user'}
+        ]
+    """
+    new_messages = []
+    for message in messages:
+        # Handle the new GPT messages format.
+        if isinstance(message, dict) and "content" in message and isinstance(message["content"], list):
+            message = copy.deepcopy(message)
+            for item in message["content"]:
+                if isinstance(item, dict) and "image_url" in item:
+                    item["image_url"]["url"] = pil_to_data_uri(item["image_url"]["url"])
+
+        new_messages.append(message)
+
+    return new_messages
+
+
+def num_tokens_from_gpt_image(
+    image_data: Union[str, Image.Image], model: str = "gpt-4-vision", low_quality: bool = False
+) -> int:
+    """
+    Calculate the number of tokens required to process an image based on its dimensions
+    after scaling for different GPT models. Supports "gpt-4-vision", "gpt-4o", and "gpt-4o-mini".
+    This function scales the image so that its longest edge is at most 2048 pixels and its shortest
+    edge is at most 768 pixels (for "gpt-4-vision"). It then calculates the number of 512x512 tiles
+    needed to cover the scaled image and computes the total tokens based on the number of these tiles.
+
+    Reference: https://openai.com/api/pricing/
+
+    Args:
+        image_data : Union[str, Image.Image]: The image data which can either be a base64
+           encoded string, a URL, a file path, or a PIL Image object.
+        model: str: The model being used for image processing. Can be "gpt-4-vision", "gpt-4o", or "gpt-4o-mini".
+
+    Returns:
+        int: The total number of tokens required for processing the image.
+
+    Examples:
+    --------
+    >>> from PIL import Image
+    >>> img = Image.new('RGB', (2500, 2500), color = 'red')
+    >>> num_tokens_from_gpt_image(img, model="gpt-4-vision")
+    765
+    """
+
+    image = get_pil_image(image_data)  # PIL Image
+    width, height = image.size
+
+    # Determine model parameters
+    if "gpt-4-vision" in model or "gpt-4-turbo" in model or "gpt-4v" in model or "gpt-4-v" in model:
+        params = MODEL_PARAMS["gpt-4-vision"]
+    elif "gpt-4o-mini" in model:
+        params = MODEL_PARAMS["gpt-4o-mini"]
+    elif "gpt-4o" in model:
+        params = MODEL_PARAMS["gpt-4o"]
+    else:
+        raise ValueError(
+            f"Model {model} is not supported. Choose 'gpt-4-vision', 'gpt-4-turbo', 'gpt-4v', 'gpt-4-v', 'gpt-4o', or 'gpt-4o-mini'."
+        )
+
+    if low_quality:
+        return params["base_token_count"]
+
+    # 1. Constrain the longest edge
+    if max(width, height) > params["max_edge"]:
+        scale_factor = params["max_edge"] / max(width, height)
+        width, height = int(width * scale_factor), int(height * scale_factor)
+
+    # 2. Further constrain the shortest edge
+    if min(width, height) > params["min_edge"]:
+        scale_factor = params["min_edge"] / min(width, height)
+        width, height = int(width * scale_factor), int(height * scale_factor)
+
+    # 3. Count how many tiles are needed to cover the image
+    tiles_width = ceil(width / params["tile_size"])
+    tiles_height = ceil(height / params["tile_size"])
+    total_tokens = params["base_token_count"] + params["token_multiplier"] * (tiles_width * tiles_height)
+
+    return total_tokens

autogen/agentchat/contrib/llamaindex_conversable_agent.py

@@ -0,0 +1,114 @@
+# 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 Dict, List, Optional, Tuple, Union
+
+from autogen import OpenAIWrapper
+from autogen.agentchat import Agent, ConversableAgent
+from autogen.agentchat.contrib.vectordb.utils import get_logger
+
+logger = get_logger(__name__)
+
+try:
+    from llama_index.core.agent.runner.base import AgentRunner
+    from llama_index.core.base.llms.types import ChatMessage
+    from llama_index.core.chat_engine.types import AgentChatResponse
+except ImportError as e:
+    logger.fatal("Failed to import llama-index. Try running 'pip install llama-index'")
+    raise e
+
+
+class LLamaIndexConversableAgent(ConversableAgent):
+    def __init__(
+        self,
+        name: str,
+        llama_index_agent: AgentRunner,
+        description: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            name (str): agent name.
+            llama_index_agent (AgentRunner): llama index agent.
+                Please override this attribute if you want to reprogram the agent.
+            description (str): a short description of the agent. This description is used by other agents
+                (e.g. the GroupChatManager) to decide when to call upon this agent.
+            **kwargs (dict): Please refer to other kwargs in
+                [ConversableAgent](../conversable_agent#__init__).
+        """
+
+        if llama_index_agent is None:
+            raise ValueError("llama_index_agent must be provided")
+
+        if description is None or description.isspace():
+            raise ValueError("description must be provided")
+
+        super().__init__(
+            name,
+            description=description,
+            **kwargs,
+        )
+
+        self._llama_index_agent = llama_index_agent
+
+        # Override the `generate_oai_reply`
+        self.replace_reply_func(ConversableAgent.generate_oai_reply, LLamaIndexConversableAgent._generate_oai_reply)
+
+        self.replace_reply_func(ConversableAgent.a_generate_oai_reply, LLamaIndexConversableAgent._a_generate_oai_reply)
+
+    def _generate_oai_reply(
+        self,
+        messages: Optional[List[Dict]] = None,
+        sender: Optional[Agent] = None,
+        config: Optional[OpenAIWrapper] = None,
+    ) -> Tuple[bool, Union[str, Dict, None]]:
+        """Generate a reply using autogen.oai."""
+        user_message, history = self._extract_message_and_history(messages=messages, sender=sender)
+
+        chatResponse: AgentChatResponse = self._llama_index_agent.chat(message=user_message, chat_history=history)
+
+        extracted_response = chatResponse.response
+
+        return (True, extracted_response)
+
+    async def _a_generate_oai_reply(
+        self,
+        messages: Optional[List[Dict]] = None,
+        sender: Optional[Agent] = None,
+        config: Optional[OpenAIWrapper] = None,
+    ) -> Tuple[bool, Union[str, Dict, None]]:
+        """Generate a reply using autogen.oai."""
+        user_message, history = self._extract_message_and_history(messages=messages, sender=sender)
+
+        chatResponse: AgentChatResponse = await self._llama_index_agent.achat(
+            message=user_message, chat_history=history
+        )
+
+        extracted_response = chatResponse.response
+
+        return (True, extracted_response)
+
+    def _extract_message_and_history(
+        self, messages: Optional[List[Dict]] = None, sender: Optional[Agent] = None
+    ) -> Tuple[str, List[ChatMessage]]:
+        """Extract the message and history from the messages."""
+        if not messages:
+            messages = self._oai_messages[sender]
+
+        if not messages:
+            return "", []
+
+        message = messages[-1].get("content", "")
+
+        history = messages[:-1]
+        history_messages: List[ChatMessage] = []
+        for history_message in history:
+            content = history_message.get("content", "")
+            role = history_message.get("role", "user")
+            if role:
+                if role == "user" or role == "assistant":
+                    history_messages.append(ChatMessage(content=content, role=role, additional_kwargs={}))
+        return message, history_messages