ag2 0.4b1__py3-none-any.whl → 0.4.1__py3-none-any.whl

autogen/agentchat/contrib/captainagent/tools/math/calculate_day_of_the_week.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def calculate_day_of_the_week(total_days: int, starting_day: str):
+    """
+    Calculates the day of the week after a given number of days starting from a specified day.
+
+    Args:
+        total_days: The number of days to calculate.
+        starting_day: The starting day of the week, should be one of 'Monday', 'Tuesday', 'Wednesday', etc.
+
+    Returns:
+        str: The day of the week after the specified number of days.
+    """
+    days_of_week = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
+
+    start_index = days_of_week.index(starting_day)
+    end_index = (start_index + total_days) % 7
+    return days_of_week[end_index]

autogen/agentchat/contrib/captainagent/tools/math/calculate_fraction_sum.py

@@ -0,0 +1,29 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def calculate_fraction_sum(
+    fraction1_numerator: int, fraction1_denominator: int, fraction2_numerator: int, fraction2_denominator: int
+):
+    """
+    Calculates the sum of two fractions and returns the result as a mixed number.
+
+    Args:
+        fraction1_numerator: The numerator of the first fraction.
+        fraction1_denominator: The denominator of the first fraction.
+        fraction2_numerator: The numerator of the second fraction.
+        fraction2_denominator: The denominator of the second fraction.
+
+    Returns:
+        str: The sum of the two fractions as a mixed number in the format 'a b/c'
+    """
+    from fractions import Fraction
+
+    fraction1 = Fraction(fraction1_numerator, fraction1_denominator)
+    fraction2 = Fraction(fraction2_numerator, fraction2_denominator)
+    result = fraction1 + fraction2
+    mixed_number = result.numerator // result.denominator
+    mixed_fraction_numerator = result.numerator % result.denominator
+    if mixed_fraction_numerator > 0:
+        return f"{mixed_number} {Fraction(mixed_fraction_numerator, result.denominator)}"
+    else:
+        return str(mixed_number)

autogen/agentchat/contrib/captainagent/tools/math/calculate_matrix_power.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+from autogen.coding.func_with_reqs import with_requirements
+
+
+@with_requirements(["sympy"])
+def calculate_matrix_power(matrix, power):
+    """
+    Calculate the power of a given matrix.
+
+    Args:
+        matrix (list): An array of numbers that represents the matrix.
+        power (int): The power to which the matrix is raised.
+
+    Returns:
+        Matrix: The resulting matrix after raising to power.
+
+    Raises:
+        ValueError: If the power is negative and the matrix is not invertible.
+    """
+    from sympy import Matrix, eye
+
+    m = Matrix(matrix)
+    if power == 0:
+        return eye(m.shape[0])
+    elif power < 0:
+        if not m.is_invertible():
+            raise ValueError("Matrix is not invertible.")
+        return m.inverse() ** (-power)
+    elif power > 0:
+        return m**power

autogen/agentchat/contrib/captainagent/tools/math/calculate_reflected_point.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def calculate_reflected_point(point):
+    """
+    Calculates the reflection point of a given point about the line y=x.
+
+    Args:
+        point (dict): A dictionary representing the coordinates of the point.
+            The dictionary should have keys 'x' and 'y' representing the x and y coordinates respectively.
+
+    Returns:
+        dict: A dictionary representing the coordinates of the reflected point. Its keys are 'x' and 'y'.
+    """
+    # Swap x and y for reflection about y=x
+    reflected_point = {"x": point["y"], "y": point["x"]}
+    return reflected_point

autogen/agentchat/contrib/captainagent/tools/math/complex_numbers_product.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+from autogen.coding.func_with_reqs import with_requirements
+
+
+@with_requirements(["sympy"])
+def complex_numbers_product(complex_numbers):
+    """
+    Calculates the product of a list of complex numbers.
+
+    Args:
+        complex_numbers (list): A list of dictionaries representing complex numbers.
+            Each dictionary should have 'real' and 'imag' keys representing the real
+            and imaginary parts of the complex number.
+
+    Returns:
+        complex: The simplified product of the complex numbers.
+
+    """
+    from sympy import I, simplify
+
+    result = 1
+    for c in complex_numbers:
+        result *= c["real"] + I * c["imag"]
+    return simplify(result)

autogen/agentchat/contrib/captainagent/tools/math/compute_currency_conversion.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+from autogen.coding.func_with_reqs import with_requirements
+
+
+@with_requirements(["sympy"])
+def compute_currency_conversion(amount, exchange_rate):
+    """
+    Compute the currency conversion of the given amount using the provided exchange rate.
+
+    Args:
+    amount (float): The amount to be converted.
+    exchange_rate (float): The exchange rate to use for the conversion, represented as the amount of second currency equivalent to one unit of the first currency.
+
+    Returns:
+    float: The converted amount.
+
+    """
+    from sympy import Rational
+
+    # Calculate the converted amount using the given exchange rate
+    converted_amount = Rational(amount, exchange_rate)
+    return float(converted_amount)

autogen/agentchat/contrib/captainagent/tools/math/count_distinct_permutations.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def count_distinct_permutations(sequence):
+    """
+    Counts the number of distinct permutations of a sequence where items may be indistinguishable.
+
+    Args:
+        sequence (iterable): The sequence for which to count the distinct permutations.
+
+    Returns:
+        int: The number of distinct permutations.
+
+    Example:
+        >>> count_distinct_permutations('aab')
+        3
+        >>> count_distinct_permutations([1, 2, 2])
+        3
+    """
+    from collections import Counter
+    from math import factorial
+
+    counts = Counter(sequence)
+    total_length = sum(counts.values())
+    permutations = factorial(total_length)
+    for count in counts.values():
+        permutations //= factorial(count)
+    return permutations

autogen/agentchat/contrib/captainagent/tools/math/evaluate_expression.py

@@ -0,0 +1,29 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def evaluate_expression(expression):
+    """
+    Evaluates a mathematical expression with support for floor function notation and power notation.
+
+    Args:
+        expression (str): The mathematical expression to evaluate. It can only contain one symbol 'x'.
+
+    Returns:
+        Union[sympy.Expr, str]: The evaluated result as a sympy expression if successful,
+        otherwise an error message as a string.
+
+    """
+    from sympy import symbols, sympify
+
+    # Replace power with ** for sympy
+    expression = expression.replace("^", "**")
+    # Replace the floor function notation
+    expression = expression.replace("\\lfloor", "floor(").replace("\\rfloor", ")")
+    try:
+        # Create a symbol 'x' for use in case it is in the expression
+        symbols("x")
+        # Evaluate the expression
+        result = sympify(expression)
+        return result
+    except Exception as e:
+        return str(e)

autogen/agentchat/contrib/captainagent/tools/math/find_continuity_point.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def find_continuity_point(f_leq, f_gt, x_value):
+    """
+    Find the value 'a' that ensures the continuity of a piecewise function at a given point.
+
+    Args:
+        f_leq (str): The function expression for f(x) when x is less than or equal to the continuity point, in the form of a string.
+        f_gt (str): The function expression for f(x) when x is greater than the continuity point, in the form of a string.
+        x_value (float): The x-value at which continuity is to be ensured.
+
+    Returns:
+        float or None: The value of 'a' that satisfies the continuity condition,
+        or None if no such value exists.
+    """
+    from sympy import Eq, solve, symbols, sympify
+
+    x, a = symbols("x a")
+
+    # Convert string to sympy expression
+    f_leq_expr = sympify(f_leq)
+    f_gt_expr = sympify(f_gt)
+
+    # Evaluate the expressions at the given x_value
+    f_leq_value = f_leq_expr.subs(x, x_value)
+    f_gt_value = f_gt_expr.subs(x, x_value)
+
+    # Set up the equation for a
+    equation = Eq(f_leq_value, f_gt_value)
+
+    # Solve the equation
+    a_value = solve(equation, a)
+
+    return a_value[0] if a_value else None

autogen/agentchat/contrib/captainagent/tools/math/fraction_to_mixed_numbers.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def fraction_to_mixed_numbers(numerator, denominator):
+    """
+    Simplifies a fraction to its lowest terms and returns it as a mixed number.
+
+    Args:
+        numerator (int): The numerator of the fraction.
+        denominator (int): The denominator of the fraction.
+
+    Returns:
+        str: The simplified fraction as a string. If the fraction is already an integer, it returns the integer as a string.
+             If the fraction is a proper fraction, it returns the mixed number representation as a string.
+             If the numerator or denominator is not an integer, it returns an error message.
+             If the denominator is zero, it returns an error message.
+    """
+    from sympy import Rational
+
+    # Ensure that numerator and denominator are integers
+    if not isinstance(numerator, int) or not isinstance(denominator, int):
+        return "Error: Numerator and denominator must be integers."
+
+    # Handle the case where the denominator is zero
+    if denominator == 0:
+        return "Error: Denominator cannot be zero."
+
+    # Simplify the fraction to its lowest terms
+    result = Rational(numerator, denominator)
+    # Return the result as a mixed number if needed
+    if result.is_integer:
+        return str(int(result))
+    else:
+        # Result as a mixed number
+        integer_part = int(result)
+        fractional_part = result - integer_part
+        if fractional_part != 0:
+            return f"{integer_part} {fractional_part}"
+        else:
+            return str(integer_part)

autogen/agentchat/contrib/captainagent/tools/math/modular_inverse_sum.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def modular_inverse_sum(expressions, modulus):
+    """
+    Calculates the sum of modular inverses of the given expressions modulo the specified modulus.
+
+    Args:
+        expressions (list): A list of numbers for which the modular inverses need to be calculated.
+        modulus (int): The modulus value.
+
+    Returns:
+        int: The sum of modular inverses modulo the specified modulus.
+    """
+    from sympy import mod_inverse
+
+    mod_sum = 0
+    for number in expressions:
+        try:
+            mod_sum += mod_inverse(number, modulus)
+        except ValueError:
+            pass  # If modular inverse does not exist, skip the term
+    return mod_sum % modulus

autogen/agentchat/contrib/captainagent/tools/math/simplify_mixed_numbers.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def simplify_mixed_numbers(numerator1, denominator1, numerator2, denominator2, whole_number1, whole_number2):
+    """
+    Simplifies the sum of two mixed numbers and returns the result as a string in the format 'a b/c'.
+
+    Args:
+        numerator1 (int): The numerator of the first fraction.
+        denominator1 (int): The denominator of the first fraction.
+        numerator2 (int): The numerator of the second fraction.
+        denominator2 (int): The denominator of the second fraction.
+        whole_number1 (int): The whole number part of the first mixed number.
+        whole_number2 (int): The whole number part of the second mixed number.
+
+    Returns:
+        str: The simplified sum of the two mixed numbers as a string in the format 'a b/c'.
+    """
+    from fractions import Fraction
+
+    # Convert mixed numbers to improper fractions
+    fraction1 = whole_number1 * denominator1 + numerator1
+    fraction2 = whole_number2 * denominator2 + numerator2
+    # Create Fraction objects
+    frac1 = Fraction(fraction1, denominator1)
+    frac2 = Fraction(fraction2, denominator2)
+    # Calculate the sum
+    result = frac1 + frac2
+    # Convert to mixed number
+    mixed_number = result.numerator // result.denominator
+    mixed_fraction_numerator = result.numerator % result.denominator
+    mixed_fraction = Fraction(mixed_fraction_numerator, result.denominator)
+    # Return as a string in the format 'a b/c'
+    if mixed_fraction_numerator > 0:
+        return f"{mixed_number} {mixed_fraction}"
+    else:
+        return str(mixed_number)

autogen/agentchat/contrib/captainagent/tools/math/sum_of_digit_factorials.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def sum_of_digit_factorials(number):
+    """
+    Calculates the sum of the factorial of each digit in a number, often used in problems involving curious numbers like 145.
+
+    Args:
+        number (int): The number for which to calculate the sum of digit factorials.
+
+    Returns:
+        int: The sum of the factorials of the digits in the given number.
+    """
+    from math import factorial
+
+    return sum(factorial(int(digit)) for digit in str(number))

autogen/agentchat/contrib/captainagent/tools/math/sum_of_primes_below.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+def sum_of_primes_below(threshold):
+    """
+    Calculates the sum of all prime numbers below a given threshold.
+
+    Args:
+        threshold (int): The maximum number (exclusive) up to which primes are summed.
+
+    Returns:
+        int: The sum of all prime numbers below the threshold.
+    """
+    from sympy import primerange
+
+    return sum(primerange(2, threshold))

autogen/agentchat/contrib/captainagent/tools/requirements.txt

@@ -0,0 +1,10 @@
+markdownify
+arxiv
+pymupdf
+wikipedia-api
+easyocr
+python-pptx
+openai-whisper
+pandas
+scipy
+sentence-transformers

autogen/agentchat/contrib/captainagent/tools/tool_description.tsv

@@ -0,0 +1,34 @@
+docid	document_content
+1	math complex_numbers_product Calculates the product of a list of complex numbers.
+2	math calculate_matrix_power Calculate the power of a given matrix.
+3	math calculate_day_of_the_week Calculates the day of the week after a given number of days starting from a specified day.
+4	math modular_inverse_sum Calculates the sum of modular inverses of the given expressions modulo the specified modulus.
+5	math sum_of_digit_factorials Calculates the sum of the factorial of each digit in a number, often used in problems involving curious numbers like 145.
+6	math sum_of_primes_below Calculates the sum of all prime numbers below a given threshold.
+7	math evaluate_expression Evaluates a mathematical expression with support for floor function notation and power notation.
+8	math compute_currency_conversion Compute the currency conversion of the given amount using the provided exchange rate.
+9	math find_continuity_point Find the value 'a' that ensures the continuity of a piecewise function at a given point.
+10	math simplify_mixed_numbers Simplifies the sum of two mixed numbers and returns the result as a string in the format 'a b/c'.
+11	math fraction_to_mixed_numbers Simplifies a fraction to its lowest terms and returns it as a mixed number.
+12	math calculate_fraction_sum Calculates the sum of two fractions and returns the result as a mixed number.
+13	math count_distinct_permutations Counts the number of distinct permutations of a sequence where items may be indistinguishable.
+14	math calculate_circle_area_from_diameter Calculate the area of a circle given its diameter.
+15	math calculate_reflected_point Calculates the reflection point of a given point about the line y=x.
+16	data_analysis explore_csv Reads a CSV file and prints the column names, shape, data types, and the first few lines of data.
+17	data_analysis calculate_correlation Calculate the correlation between two columns in a CSV file.
+18	data_analysis detect_outlier_zscore Detect outliers in a CSV file based on a specified column. The outliers are determined by calculating the z-score of the data points in the column.
+19	data_analysis detect_outlier_iqr Detect outliers in a specified column of a CSV file using the IQR method.
+20	data_analysis shapiro_wilk_test Perform the Shapiro-Wilk test on a specified column of a CSV file.
+21	data_analysis calculate_skewness_and_kurtosis Calculate the skewness and kurtosis of a specified column in a CSV file. The kurtosis is calculated using the Fisher definition. The two metrics are computed using scipy.stats functions.
+22	information_retrieval perform_web_search Perform a web search using Bing API.
+23	information_retrieval transcribe_audio_file Transcribes the audio file located at the given file path.
+24	information_retrieval arxiv_search Search for articles on arXiv based on the given query.
+25	information_retrieval arxiv_download Downloads PDF files from ArXiv based on a list of arxiv paper IDs.
+26	information_retrieval scrape_wikipedia_tables Scrapes Wikipedia tables based on a given URL and header keyword.
+27	information_retrieval extract_pdf_text Extracts text from a specified page or the entire PDF file.
+28	information_retrieval extract_pdf_image Extracts images from a PDF file and saves them to the specified output directory.
+29	information_retrieval image_qa Perform question answering on an image using a pre-trained VQA model.
+30	information_retrieval optical_character_recognition Perform optical character recognition (OCR) on the given image.
+31	information_retrieval get_youtube_caption Retrieves the captions for a YouTube video.
+32	information_retrieval youtube_download Downloads a YouTube video and returns the download link.
+33	information_retrieval get_wikipedia_text Retrieves the text content of a Wikipedia page. It does not support tables and other complex formatting.

autogen/agentchat/contrib/captainagent.py

@@ -1,3 +1,6 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
 import hashlib
 import json
 import os

autogen/agentchat/contrib/graph_rag/document.py

@@ -1,3 +1,9 @@
+# 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 dataclasses import dataclass
 from enum import Enum, auto
 from typing import Optional

autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py

@@ -1,21 +1,24 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
 import os
-from dataclasses import field
+from dataclasses import dataclass, field
 from typing import List
 
 from graphrag_sdk import KnowledgeGraph, Source
-from graphrag_sdk.schema import Schema
+from graphrag_sdk.model_config import KnowledgeGraphModelConfig
+from graphrag_sdk.models import GenerativeModel
+from graphrag_sdk.models.openai import OpenAiGenerativeModel
+from graphrag_sdk.ontology import Ontology
 
 from .document import Document
 from .graph_query_engine import GraphStoreQueryResult
 
 
-class FalkorGraphQueryResult(GraphStoreQueryResult):
-    messages: list = field(default_factory=list)
-
-
 class FalkorGraphQueryEngine:
     """
-    This is a wrapper for Falkor DB KnowledgeGraph.
+    This is a wrapper for FalkorDB KnowledgeGraph.
     """
 
     def __init__(
@@ -25,11 +28,11 @@ class FalkorGraphQueryEngine:
         port: int = 6379,
         username: str | None = None,
         password: str | None = None,
-        model: str = "gpt-4-1106-preview",
-        schema: Schema | None = None,
+        model: GenerativeModel = OpenAiGenerativeModel("gpt-4o"),
+        ontology: Ontology | None = None,
     ):
         """
-        Initialize a Falkor DB knowledge graph.
+        Initialize a FalkorDB knowledge graph.
         Please also refer to https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/kg.py
 
         Args:
@@ -38,11 +41,18 @@ class FalkorGraphQueryEngine:
             port (int): FalkorDB port number.
             username (str|None): FalkorDB username.
             password (str|None): FalkorDB password.
-            model (str): OpenAI model to use for Falkor DB to build and retrieve from the graph.
-            schema: Falkor DB knowledge graph schema (ontology), https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/schema/schema.py
-                    If None, Falkor DB will auto generate a schema from the input docs.
+            model (GenerativeModel): LLM model to use for FalkorDB to build and retrieve from the graph, default to use OAI gpt-4o.
+            ontology: FalkorDB knowledge graph schema/ontology, https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/ontology.py
+                If None, FalkorDB will auto generate an ontology from the input docs.
         """
-        self.knowledge_graph = KnowledgeGraph(name, host, port, username, password, model, schema)
+        self.name = name
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.model = model
+        self.model_config = KnowledgeGraphModelConfig.with_model(model)
+        self.ontology = ontology
 
     def init_db(self, input_doc: List[Document] | None):
         """
@@ -54,14 +64,33 @@ class FalkorGraphQueryEngine:
                 sources.append(Source(doc.path_or_url))
 
         if sources:
+            # Auto generate graph ontology if not created by user.
+            if self.ontology is None:
+                self.ontology = Ontology.from_sources(
+                    sources=sources,
+                    model=self.model,
+                )
+
+            self.knowledge_graph = KnowledgeGraph(
+                name=self.name,
+                host=self.host,
+                port=self.port,
+                username=self.username,
+                password=self.password,
+                model_config=KnowledgeGraphModelConfig.with_model(self.model),
+                ontology=self.ontology,
+            )
+
+            # Establish a chat session, this will maintain the history
+            self._chat_session = self.knowledge_graph.chat_session()
             self.knowledge_graph.process_sources(sources)
 
     def add_records(self, new_records: List) -> bool:
-        raise NotImplementedError("This method is not supported by Falkor DB SDK yet.")
+        raise NotImplementedError("This method is not supported by FalkorDB SDK yet.")
 
-    def query(self, question: str, n_results: int = 1, **kwargs) -> FalkorGraphQueryResult:
+    def query(self, question: str, n_results: int = 1, **kwargs) -> GraphStoreQueryResult:
         """
-        Query the knowledage graph with a question and optional message history.
+        Query the knowledge graph with a question and optional message history.
 
         Args:
         question: a human input question.
@@ -71,6 +100,12 @@ class FalkorGraphQueryEngine:
 
         Returns: FalkorGraphQueryResult
         """
-        messages = kwargs.pop("messages", [])
-        answer, messages = self.knowledge_graph.ask(question, messages)
-        return FalkorGraphQueryResult(answer=answer, results=[], messages=messages)
+        if self.knowledge_graph is None:
+            raise ValueError("Knowledge graph is not created.")
+
+        response = self._chat_session.send_message(question)
+
+        # History will be considered when querying by setting the last_answer
+        self._chat_session.last_answer = response["response"]
+
+        return GraphStoreQueryResult(answer=response["response"], results=[])

autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from autogen import Agent, ConversableAgent, UserProxyAgent
+
+from .falkor_graph_query_engine import FalkorGraphQueryEngine
+from .graph_query_engine import GraphStoreQueryResult
+from .graph_rag_capability import GraphRagCapability
+
+
+class FalkorGraphRagCapability(GraphRagCapability):
+    """
+    The FalkorDB GraphRAG capability integrate FalkorDB with graphrag_sdk version: 0.1.3b0.
+    Ref: https://github.com/FalkorDB/GraphRAG-SDK/tree/2-move-away-from-sql-to-json-ontology-detection
+
+    For usage, please refer to example notebook/agentchat_graph_rag_falkordb.ipynb
+    """
+
+    def __init__(self, query_engine: FalkorGraphQueryEngine):
+        """
+        initialize GraphRAG capability with a graph query engine
+        """
+        self.query_engine = query_engine
+
+    def add_to_agent(self, agent: UserProxyAgent):
+        """
+        Add FalkorDB GraphRAG capability to a UserProxyAgent.
+        The restriction to a UserProxyAgent to make sure the returned message does not contain information retrieved from the graph DB instead of any LLMs.
+        """
+        self.graph_rag_agent = agent
+
+        # Validate the agent config
+        if agent.llm_config not in (None, False):
+            raise Exception(
+                "Agents with GraphRAG capabilities do not use an LLM configuration. Please set your llm_config to None or False."
+            )
+
+        # Register method to generate the reply using a FalkorDB query
+        # All other reply methods will be removed
+        agent.register_reply(
+            [ConversableAgent, None], self._reply_using_falkordb_query, position=0, remove_other_reply_funcs=True
+        )
+
+    def _reply_using_falkordb_query(
+        self,
+        recipient: ConversableAgent,
+        messages: Optional[List[Dict]] = None,
+        sender: Optional[Agent] = None,
+        config: Optional[Any] = None,
+    ) -> Tuple[bool, Union[str, Dict, None]]:
+        """
+        Query FalkorDB and return the message. Internally, it utilises OpenAI to generate a reply based on the given messages.
+        The history with FalkorDB is also logged and updated.
+
+        If no results are found, a default message is returned: "I'm sorry, I don't have an answer for that."
+
+        Args:
+            recipient: The agent instance that will receive the message.
+            messages: A list of messages in the conversation history with the sender.
+            sender: The agent instance that sent the message.
+            config: Optional configuration for message processing.
+
+        Returns:
+            A tuple containing a boolean indicating success and the assistant's reply.
+        """
+        question = self._get_last_question(messages[-1])
+        result: GraphStoreQueryResult = self.query_engine.query(question)
+
+        return True, result.answer if result.answer else "I'm sorry, I don't have an answer for that."
+
+    def _get_last_question(self, message: Union[Dict, str]):
+        """Retrieves the last message from the conversation history."""
+        if isinstance(message, str):
+            return message
+        if isinstance(message, Dict):
+            if "content" in message:
+                return message["content"]
+        return None