graphrag_sdk 0.8.1__py3-none-any.whl

graphrag_sdk/__init__.py

@@ -0,0 +1,40 @@
+from .source import Source
+from .ontology import Ontology
+from .kg import KnowledgeGraph
+from .model_config import KnowledgeGraphModelConfig
+from .steps.create_ontology_step import CreateOntologyStep
+from .models.model import (
+    GenerativeModel,
+    GenerationResponse,
+    GenerativeModelChatSession,
+    GenerativeModelConfig,
+    FinishReason,
+)
+from .entity import Entity
+from .relation import Relation
+from .attribute import Attribute, AttributeType
+
+# Setup Null handler
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+logger.addHandler(logging.NullHandler())
+
+
+__all__ = [
+    "Source",
+    "Ontology",
+    "KnowledgeGraph",
+    "KnowledgeGraphModelConfig",
+    "CreateOntologyStep",
+    "GenerativeModel",
+    "GenerationResponse",
+    "GenerativeModelChatSession",
+    "GenerativeModelConfig",
+    "FinishReason",
+    "Entity",
+    "Relation",
+    "Attribute",
+    "AttributeType",
+]

graphrag_sdk/agents/__init__.py

@@ -0,0 +1,3 @@
+from .agent import Agent
+
+__all__ = ['Agent']

graphrag_sdk/agents/agent.py

@@ -0,0 +1,165 @@
+from abc import ABC, abstractmethod
+
+class AgentResponseCode:
+    """
+    Represents the response codes for an agent.
+    """
+
+    AGENT_RESPONSE = "agent_response"
+    AGENT_ERROR = "agent_error"
+    AGENT_REQUEST_INPUT = "agent_request_input"
+
+    @staticmethod
+    def from_str(text: str) -> "AgentResponseCode":
+        """
+        Converts a string to an AgentResponseCode object.
+
+        Args:
+            text (str): The string representation of the response code.
+
+        Returns:
+            AgentResponseCode: The corresponding AgentResponseCode object.
+
+        Raises:
+            ValueError: If the input string does not match any known response code.
+        """
+        if text == AgentResponseCode.AGENT_RESPONSE:
+            return AgentResponseCode.AGENT_RESPONSE
+        elif text == AgentResponseCode.AGENT_ERROR:
+            return AgentResponseCode.AGENT_ERROR
+        elif text == AgentResponseCode.AGENT_REQUEST_INPUT:
+            return AgentResponseCode.AGENT_REQUEST_INPUT
+        else:
+            raise ValueError(f"Unknown agent response code: {text}")
+
+
+class AgentResponse:
+    """
+    Represents a response from an agent.
+
+    Attributes:
+        response_code (AgentResponseCode): The response code.
+        payload (dict): The payload of the response.
+    """
+
+    def __init__(self, response_code: AgentResponseCode, payload: dict):
+        """
+        Initializes a new Agent object.
+
+        Args:
+            response_code (AgentResponseCode): The response code of the agent.
+            payload (dict): The payload associated with the agent.
+
+        """
+        self.response_code = response_code
+        self.payload = payload
+
+    def to_json(self) -> dict:
+        """
+        Converts the AgentResponse object to a JSON-compatible dictionary.
+
+        Returns:
+            dict: The JSON representation of the AgentResponse object.
+        """
+        return {
+            "response_code": self.response_code,
+            "payload": self.payload,
+        }
+
+    @staticmethod
+    def from_json(json: dict) -> "AgentResponse":
+        """
+        Creates an AgentResponse object from a JSON-compatible dictionary.
+
+        Args:
+            json (dict): The JSON representation of the AgentResponse object.
+
+        Returns:
+            AgentResponse: The created AgentResponse object.
+        """
+        return AgentResponse(
+            AgentResponseCode.from_str(json["response_code"]),
+            json["payload"],
+        )
+
+    def __str__(self) -> str:
+        """
+        Returns a string representation of the AgentResponse object.
+
+        Returns:
+            str: The string representation of the AgentResponse object.
+        """
+        return (
+            f"AgentResponse(response_code={self.response_code}, payload={self.payload})"
+        )
+
+    def __repr__(self) -> str:
+        """
+        Returns a string representation of the AgentResponse object.
+
+        Returns:
+            str: The string representation of the AgentResponse object.
+        """
+        return str(self)
+
+
+class Agent(ABC):
+    """
+    Abstract base class for agents in the system.
+    """
+
+    @property
+    @abstractmethod
+    def agent_id(self) -> str:
+        """
+        Get the unique identifier of the agent.
+
+        Returns:
+            str: The agent's identifier.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def introduction(self) -> str:
+        """
+        Get the introduction of the agent.
+
+        Returns:
+            str: The agent's introduction.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def interface(self) -> list[dict]:
+        """
+        Get the interface of the agent.
+
+        Returns:
+            List[Dict]: The agent's interface.
+        """
+        pass
+
+    @abstractmethod
+    def run(self, params: dict) -> str:
+        """
+        Run the agent with the given parameters and chat session.
+
+        Args:
+            params (Dict): The parameters for the agent.
+
+        Returns:
+            str: The agent's response.
+        """
+        pass
+
+    @abstractmethod
+    def __repr__(self) -> str:
+        """
+        Get a string representation of the agent.
+
+        Returns:
+            str: The string representation of the agent.
+        """
+        pass

graphrag_sdk/agents/kg_agent.py

@@ -0,0 +1,148 @@
+from .agent import Agent
+from graphrag_sdk.kg import KnowledgeGraph
+
+
+class KGAgent(Agent):
+    """Represents an Agent for a FalkorDB Knowledge Graph.
+
+    Args:
+        agent_id (str): The ID of the agent.
+        kg (KnowledgeGraph): The knowledge graph to query.
+        introduction (str): The introduction to the agent.
+
+    Examples:
+        >>> from graphrag_sdk import KnowledgeGraph, Orchestrator
+        >>> from graphrag_sdk.agents.kg_agent import KGAgent
+        >>> orchestrator = Orchestrator(model)
+        >>> kg = KnowledgeGraph("test_kg", ontology, model)
+        >>> agent = KGAgent("test_agent", kg, "This is a test agent.")
+        >>> orchestrator.register_agent(agent)
+
+    """
+
+    _interface = [
+        {
+            "name": "prompt",
+            "type": "string",
+            "required": True,
+            "description": "The prompt to ask the agent.",
+        }
+    ]
+
+    def __init__(self, agent_id: str, kg: KnowledgeGraph, introduction: str):
+        """
+        Initializes a KGAgent object.
+
+        Args:
+            agent_id (str): The ID of the agent.
+            kg (KnowledgeGraph): The knowledge graph associated with the agent.
+            introduction (str): The introduction of the agent.
+        """
+        super().__init__()
+        self.agent_id = agent_id
+        self.introduction = introduction
+        self.kg = kg
+        self.chat_session = self._kg.chat_session()
+
+    @property
+    def agent_id(self) -> str:
+        """
+        Returns the ID of the agent.
+
+        Returns:
+            str: The ID of the agent.
+        """
+        return self._agent_id
+
+    @agent_id.setter
+    def agent_id(self, value) -> None:
+        """
+        Sets the agent ID.
+
+        Args:
+            value (str): The ID of the agent.
+        """
+        self._agent_id = value
+
+    @property
+    def introduction(self) -> str:
+        """
+        Returns the introduction of the agent.
+
+        Returns:
+            str: The introduction of the agent.
+        """
+        return self._introduction
+
+    @introduction.setter
+    def introduction(self, value) -> None:
+        """
+        Sets the introduction of the agent.
+
+        Args:
+            value (str): The introduction of the agent.
+        """
+        self._introduction = value
+
+    @property
+    def interface(self) -> list[dict]:
+        """
+        Returns the interface of the KG agent.
+
+        Returns:
+            list[dict]: The interface of the KG agent.
+        """
+        return self._interface
+
+    @property
+    def kg(self) -> KnowledgeGraph:
+        """
+        Returns the KnowledgeGraph object associated with this agent.
+
+        Returns:
+            KnowledgeGraph: The KnowledgeGraph object.
+        """
+        return self._kg
+
+    @kg.setter
+    def kg(self, value: KnowledgeGraph) -> None:
+        """
+        Sets the knowledge graph for the agent.
+
+        Args:
+            value (KnowledgeGraph): The knowledge graph to be set.
+        """
+        self._kg = value
+
+    def run(self, params: dict) -> str:
+        """
+        Ask the agent a question.
+
+        Args:
+            params (dict): The parameters for the agent.
+
+        Returns:
+            str: The agent's response.
+
+        """
+        output = self.chat_session.send_message(params["prompt"])
+        return output['response']
+
+    def __repr__(self) -> str:
+        """
+        Returns a string representation of the KGAgent object.
+
+        The string representation includes the Agent ID, Knowledge Graph Name,
+        Interface, and Introduction.
+
+        Returns:
+            str: A string representation of the KGAgent object.
+        """
+        return f"""
+---
+Agent ID: {self.agent_id}
+Knowledge Graph Name: {self._kg.name}
+Interface: {self.interface}
+
+Introduction: {self.introduction}
+"""

graphrag_sdk/attribute.py

@@ -0,0 +1,170 @@
+import re
+import json
+import logging
+from enum import Enum
+from typing import Union, Optional
+
+
+logger = logging.getLogger(__name__)
+
+class AttributeType(Enum):
+    """
+    Represents the types of attributes in the system.
+    """
+
+    STRING = "string"
+    NUMBER = "number"
+    BOOLEAN = "boolean"
+    LIST = "list"
+    POINT = "point"
+    MAP = "map"
+    VECTOR = "vectorf32"
+
+    @staticmethod
+    def from_string(txt: str) -> "AttributeType":
+        """
+        Converts a string representation of an attribute type to its corresponding AttributeType value.
+
+        Args:
+            txt (str): The string representation of the attribute type.
+
+        Returns:
+            str: The corresponding AttributeType value.
+
+        Raises:
+            ValueError: If the provided attribute type is invalid.
+        """
+        # Graph representation of the attribute type
+        normalized_txt = txt.lstrip('AttributeType.').lower()
+        
+        # Find the matching attribute type
+        if normalized_txt in _SYNONYMS:
+            return _SYNONYMS[normalized_txt]
+        
+        raise ValueError(f"Invalid attribute type: {txt}")
+    
+# Mapping of string representations to AttributeType enum members.
+_SYNONYMS = {
+    "string": AttributeType.STRING,
+    "integer": AttributeType.NUMBER,
+    "float": AttributeType.NUMBER,
+    "number": AttributeType.NUMBER,
+    "boolean": AttributeType.BOOLEAN,
+    "list": AttributeType.LIST,
+    "point": AttributeType.POINT,
+    "map": AttributeType.MAP,
+    "vectorf32": AttributeType.VECTOR,
+}
+
+class Attribute:
+    """ Represents an attribute of an entity or relation in the ontology.
+
+        Args:
+            name (str): The name of the attribute.
+            attr_type (AttributeType): The type of the attribute.
+            unique (bool): Whether the attribute is unique.
+            required (bool): Whether the attribute is required.
+
+        Examples:
+            >>> attr = Attribute("name", AttributeType.STRING, True, True)
+            >>> print(attr)
+            name: "string!*"
+    """
+
+    def __init__(
+        self, name: str, attr_type: AttributeType, unique: Optional[bool] = False, required: Optional[bool] = False
+    ):
+        """
+        Initialize a new Attribute object.
+
+        Args:
+            name (str): The name of the attribute.
+            attr_type (AttributeType): The type of the attribute.
+            unique (Optional[bool]): Indicates whether the attribute should be unique. Defaults to False.
+            required (Optional[bool]): Indicates whether the attribute is required. Defaults to False.
+        """
+        self.name = re.sub(r"([^a-zA-Z0-9_])", "_", name)
+        self.type = attr_type
+        self.unique = unique
+        self.required = required
+
+    @staticmethod
+    def from_json(txt: Union[str, dict]) -> "Attribute":
+        """
+        Creates an Attribute object from a JSON string or dictionary.
+
+        Args:
+            txt (Union[str, dict]): The JSON string or dictionary representing the Attribute.
+
+        Returns:
+            Attribute: The created Attribute object.
+
+        """
+        txt = txt if isinstance(txt, dict) else json.loads(txt)
+
+        return Attribute(
+            txt["name"],
+            AttributeType.from_string(txt["type"]),
+            txt["unique"],
+            txt["required"] if "required" in txt else False,
+        )
+
+    @staticmethod
+    def from_string(txt: str) -> "Attribute":
+        """
+        Parses an attribute from a string.
+        The "!" symbol indicates that the attribute is unique.
+        The "*" symbol indicates that the attribute is required
+
+        Args:
+            txt (str): The string to parse.
+
+        Returns:
+            Attribute: The parsed attribute.
+
+        Examples:
+            >>> attr = Attribute.from_string("name:string!*")
+            >>> print(attr.name)
+            name
+
+        Raises:
+            Exception: If the attribute type is invalid.
+        """
+        name = txt.split(":")[0].strip()
+        attr_type = txt.split(":")[1].split("!")[0].split("*")[0].strip()
+        unique = "!" in txt
+        required = "*" in txt
+
+        return Attribute(name, AttributeType.from_string(attr_type), unique, required)
+
+    def to_json(self) -> dict:
+        """
+        Converts the attribute object to a JSON representation.
+
+        Returns:
+            dict: A dictionary representing the attribute object in JSON format.
+                The dictionary contains the following keys:
+                - "name": The name of the attribute.
+                - "type": The type of the attribute.
+                - "unique": A boolean indicating whether the attribute is unique.
+                - "required": A boolean indicating whether the attribute is required.
+        """
+        json_data = {
+            "name": self.name,
+            "type": self.type.value,
+            "unique": self.unique,
+            "required": self.required,
+        }
+
+        return json_data
+
+    def __str__(self) -> str:
+        """
+        Returns a string representation of the Attribute object.
+
+        The string representation includes the attribute name, type, uniqueness, and requirement status.
+
+        Returns:
+            str: A string representation of the Attribute object.
+        """
+        return f"{self.name}: \"{self.type}{'!' if self.unique else ''}{'*' if self.required else ''}\""