gllm-datastore-binary 0.5.45__cp311-cp311-macosx_13_0_arm64.whl

gllm_datastore/graph_data_store/light_rag_data_store.pyi

@@ -0,0 +1,93 @@
+from _typeshed import Incomplete
+from gllm_core.schema import Chunk as Chunk
+from gllm_datastore.graph_data_store.graph_rag_data_store import BaseGraphRAGDataStore as BaseGraphRAGDataStore
+from gllm_datastore.graph_data_store.utils.constants import LightRAGConstants as LightRAGConstants, LightRAGKeys as LightRAGKeys
+from gllm_datastore.graph_data_store.utils.light_rag_em_invoker_adapter import LightRAGEMInvokerAdapter as LightRAGEMInvokerAdapter
+from lightrag import LightRAG
+from typing import Any
+
+class BaseLightRAGDataStore(BaseGraphRAGDataStore):
+    """LightRAG data store base class.
+
+    This class provides an abstract base class for the BaseGraphRAGDataStore interface
+    using LightRAG as the underlying technology. It handles indexing files
+    into a graph database, creating relationships between files and chunks,
+    and provides methods for deleting files and chunks from the graph.
+
+    Please use LightRAGPostgresDataStore or other concrete implementations instead.
+
+    To implement a concrete data store, inherit from this class and implement the
+    abstract methods.
+
+    Attributes:
+        instance (LightRAG): The LightRAG instance to use.
+        is_initialized (bool): Whether the data store is initialized.
+    """
+    instance: Incomplete
+    is_initialized: bool
+    def __init__(self, instance: LightRAG) -> None:
+        """Initialize the LightRAG data store.
+
+        This is an abstract base class and cannot be instantiated directly.
+        Use LightRAGPostgresDataStore or other concrete implementations instead.
+
+        Args:
+            instance (LightRAG): The LightRAG instance to use for indexing.
+
+        Raises:
+            TypeError: If attempting to instantiate BaseLightRAGDataStore directly.
+        """
+    async def ensure_initialized(self) -> None:
+        """Ensure that the LightRAG data store is initialized.
+
+        This asynchronous method ensures that the LightRAG data store is initialized.
+        If the data store is not initialized, it will initialize it.
+        """
+    async def map_file_id_to_chunk_ids_using_graph(self, file_id: str, chunk_ids: list[str]) -> None:
+        """Create file and chunk nodes in the graph and establish relationships.
+
+        This asynchronous method creates a file node and multiple chunk nodes
+        in the graph database, then establishes relationships between the file
+        and its chunks. The relationships are necessary for retaining the relationships
+        between the file and its chunks when deleting the file.
+
+        Args:
+            file_id (str): The ID of the file to create in the graph.
+            chunk_ids (list[str]): List of chunk IDs that belong to the file.
+        """
+    async def insert(self, chunks: list[Chunk]) -> None:
+        """Insert a chunk into the LightRAG data store.
+
+        This asynchronous method inserts a chunk into the LightRAG data store.
+        If the data store is not initialized, it will initialize it.
+
+        Args:
+            chunks (list[Chunk]): The chunks to insert.
+        """
+    async def query(self, query: str, **kwargs: Any) -> Any:
+        """Query the LightRAG data store.
+
+        Args:
+            query (str): The query to be executed.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            Any: The result of the query.
+        """
+    async def delete(self, chunk_id: str) -> None:
+        """Delete a chunk from the LightRAG data store.
+
+        Args:
+            chunk_id (str): The ID of the chunk to delete.
+        """
+    async def delete_by_document_id(self, document_id: str, **kwargs: Any) -> None:
+        """Delete a document/file and all its associated chunks from the LightRAG data store.
+
+        This asynchronous method retrieves all chunks associated with a document/file,
+        deletes each chunk from both the LightRAG system and the graph database,
+        and finally deletes the document/file node itself.
+
+        Args:
+            document_id (str): The ID of the document to delete.
+            **kwargs (Any): Additional keyword arguments.
+        """

gllm_datastore/graph_data_store/light_rag_postgres_data_store.pyi

@@ -0,0 +1,96 @@
+from _typeshed import Incomplete
+from contextlib import contextmanager
+from gllm_datastore.graph_data_store.light_rag_data_store import BaseLightRAGDataStore as BaseLightRAGDataStore
+from gllm_datastore.graph_data_store.utils.constants import LightRAGPostgresStorageConstants as LightRAGPostgresStorageConstants
+from gllm_datastore.graph_data_store.utils.light_rag_em_invoker_adapter import LightRAGEMInvokerAdapter as LightRAGEMInvokerAdapter
+from gllm_datastore.graph_data_store.utils.light_rag_lm_invoker_adapter import LightRAGLMInvokerAdapter as LightRAGLMInvokerAdapter
+from gllm_inference.em_invoker.em_invoker import BaseEMInvoker
+from gllm_inference.lm_invoker.lm_invoker import BaseLMInvoker
+from lightrag import LightRAG
+from pydantic import BaseModel
+from typing import Any, Generator
+
+class PostgresDBConfig(BaseModel):
+    """Pydantic model containing PostgreSQL configuration parameters."""
+    host: str
+    port: int
+    user: str
+    password: str
+    database: str
+    workspace: str
+
+@contextmanager
+def postgres_config_context(config: PostgresDBConfig) -> Generator[None, None, None]:
+    """Context manager to temporarily set postgres config.
+
+    This context manager is used to temporarily set the environment variables
+    for PostgreSQL configuration. It is used to set the environment variables
+    for the duration of the context manager's block as it is required by LightRAG.
+
+    Args:
+        config (PostgresDBConfig): Pydantic model containing PostgreSQL configuration parameters.
+
+    Returns:
+        Generator[None, None, None]: Generator object that can be used in a context manager.
+    """
+
+class LightRAGPostgresDataStore(BaseLightRAGDataStore):
+    '''Data store implementation for LightRAG-based graph RAG using PostgreSQL.
+
+    This class extends the LightRAGDataStore to use PostgreSQL as the graph database,
+    key-value store, and vector database.
+
+    To use this data store, please ensure that you have a PostgreSQL with AGE and PGVector extensions installed.
+    You can use the following docker run command to start a PostgreSQL container with AGE and PGVector extensions:
+
+    ```bash
+    docker run         -p 5455:5432         -d         --name postgres-LightRag         shangor/postgres-for-rag:v1.0         sh -c "service postgresql start && sleep infinity"
+    ```
+
+    Example:
+        ```python
+        from gllm_inference.em_invoker import OpenAIEMInvoker
+        from gllm_inference.lm_invoker import OpenAILMInvoker
+        from gllm_datastore.graph_data_store.light_rag_postgres_data_store import LightRAGPostgresDataStore
+
+        # Create the indexer
+        data_store = await LightRAGPostgresDataStore(
+            lm_invoker=OpenAILMInvoker(model_name="gpt-4o-mini"),
+            em_invoker=OpenAIEMInvoker(model_name="text-embedding-3-small"),
+            postgres_db_user="rag",
+            postgres_db_password="rag",
+            postgres_db_name="rag",
+            postgres_db_host="localhost",
+            postgres_db_port=5455,
+        )
+
+        # Retrieve using LightRAG instance
+        await data_store.query("What is AI?")
+        ```
+
+    Attributes:
+        instance (LightRAG): The LightRAG instance used for indexing and querying.
+        lm_invoker_adapter (LightRAGLMInvokerAdapter): The adapter for the LM invoker.
+        em_invoker_adapter (LightRAGEMInvokerAdapter): The adapter for the EM invoker.
+        postgres_config (PostgresDBConfig): Pydantic model containing PostgreSQL configuration parameters.
+    '''
+    lm_invoker_adapter: Incomplete
+    em_invoker_adapter: Incomplete
+    postgres_config: Incomplete
+    def __init__(self, lm_invoker: BaseLMInvoker, em_invoker: BaseEMInvoker, postgres_db_host: str = 'localhost', postgres_db_port: int = 5432, postgres_db_user: str = 'postgres', postgres_db_password: str = 'password', postgres_db_name: str = 'postgres', postgres_db_workspace: str = 'default', use_cache: bool = False, lm_invoke_kwargs: dict[str, Any] | None = None, instance: LightRAG | None = None, **kwargs: Any) -> None:
+        '''Initialize the LightRAGPostgresIndexer.
+
+        Args:
+            lm_invoker (BaseLMInvoker): The LM invoker to use.
+            em_invoker (BaseEMInvoker): The EM invoker to use.
+            postgres_db_host (str, optional): The host for the PostgreSQL database. Defaults to "localhost".
+            postgres_db_port (int, optional): The port for the PostgreSQL database. Defaults to 5432.
+            postgres_db_user (str, optional): The user for the PostgreSQL database. Defaults to "postgres".
+            postgres_db_password (str, optional): The password for the PostgreSQL database. Defaults to "password".
+            postgres_db_name (str, optional): The name for the PostgreSQL database. Defaults to "postgres".
+            postgres_db_workspace (str, optional): The workspace for the PostgreSQL database. Defaults to "default".
+            use_cache (bool, optional): Whether to enable caching for the LightRAG instance. Defaults to False.
+            lm_invoke_kwargs (dict[str, Any] | None, optional): Keyword arguments for the LM invoker. Defaults to None.
+            instance (LightRAG | None, optional): A configured LightRAG instance to use. Defaults to None.
+            **kwargs (Any): Additional keyword arguments.
+        '''

gllm_datastore/graph_data_store/llama_index_graph_rag_data_store.pyi

@@ -0,0 +1,49 @@
+from _typeshed import Incomplete
+from abc import ABC
+from gllm_datastore.graph_data_store.graph_rag_data_store import BaseGraphRAGDataStore as BaseGraphRAGDataStore
+from gllm_datastore.graph_data_store.utils import LlamaIndexEMInvokerAdapter as LlamaIndexEMInvokerAdapter, LlamaIndexLMInvokerAdapter as LlamaIndexLMInvokerAdapter
+from gllm_inference.em_invoker.em_invoker import BaseEMInvoker
+from gllm_inference.lm_invoker.lm_invoker import BaseLMInvoker
+from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.graph_stores.types import PropertyGraphStore
+from llama_index.core.llms import LLM
+from typing import Any
+
+class LlamaIndexGraphRAGDataStore(PropertyGraphStore, BaseGraphRAGDataStore, ABC):
+    """Abstract base class for a LlamaIndex graph RAG data store.
+
+    This class provides a common interface for LlamaIndex-based graph RAG data stores.
+    It handles conversion from GLLM invokers to LlamaIndex LLM and embedding models.
+
+    Attributes:
+        lm_invoker (BaseLMInvoker | None): The GLLM language model invoker.
+        em_invoker (BaseEMInvoker | None): The GLLM embedding model invoker.
+        llm (LLM | None): The LlamaIndex LLM instance (converted from lm_invoker if provided).
+        embed_model (BaseEmbedding | None): The LlamaIndex embedding instance (converted from em_invoker if provided).
+    """
+    lm_invoker: Incomplete
+    em_invoker: Incomplete
+    llm: LLM | None
+    embed_model: BaseEmbedding | None
+    def __init__(self, lm_invoker: BaseLMInvoker | None = None, em_invoker: BaseEMInvoker | None = None, **kwargs: Any) -> None:
+        """Initialize the LlamaIndexGraphRAGDataStore.
+
+        Args:
+            lm_invoker (BaseLMInvoker | None, optional): GLLM language model invoker.
+                If provided, it will be automatically converted to a LlamaIndex LLM instance
+                using LlamaIndexLMInvokerAdapter. Defaults to None.
+            em_invoker (BaseEMInvoker | None, optional): GLLM embedding model invoker.
+                If provided, it will be automatically converted to a LlamaIndex BaseEmbedding instance
+                using LlamaIndexEMInvokerAdapter. Defaults to None.
+            **kwargs (Any): Additional keyword arguments passed to PropertyGraphStore.
+        """
+    async def query(self, query: str, **kwargs: Any) -> Any:
+        """Query the graph RAG data store.
+
+        Args:
+            query (str): The query to be executed.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            Any: The result of the query.
+        """

gllm_datastore/graph_data_store/llama_index_neo4j_graph_rag_data_store.pyi

@@ -0,0 +1,78 @@
+from _typeshed import Incomplete
+from gllm_datastore.graph_data_store.llama_index_graph_rag_data_store import LlamaIndexGraphRAGDataStore as LlamaIndexGraphRAGDataStore
+from gllm_inference.em_invoker.em_invoker import BaseEMInvoker
+from gllm_inference.lm_invoker.lm_invoker import BaseLMInvoker
+from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore
+from typing import Any
+
+class LlamaIndexNeo4jGraphRAGDataStore(LlamaIndexGraphRAGDataStore, Neo4jPropertyGraphStore):
+    '''Graph RAG data store for Neo4j.
+
+    This class extends the Neo4jPropertyGraphStore class from LlamaIndex.
+    This class provides an interface for graph-based Retrieval-Augmented Generation (RAG)
+    operations on Neo4j graph databases.
+
+    Attributes:
+        neo4j_version_tuple (tuple[int, ...]): The Neo4j version tuple.
+        lm_invoker (BaseLMInvoker | None): The GLLM language model invoker (inherited from parent).
+        em_invoker (BaseEMInvoker | None): The GLLM embedding model invoker (inherited from parent).
+        llm (LLM | None): The LlamaIndex LLM instance (converted from lm_invoker, inherited from parent).
+        embed_model (BaseEmbedding | None): The LlamaIndex embedding instance
+            (converted from em_invoker, inherited from parent).
+
+    Example:
+        ```python
+        # Option 1: Use with GLLM invokers (recommended)
+        from gllm_inference.builder import build_lm_invoker, build_em_invoker
+
+        lm_invoker = build_lm_invoker(model_id="openai/gpt-4o-mini")
+        em_invoker = build_em_invoker(model_id="openai/text-embedding-3-small")
+
+        store = LlamaIndexNeo4jGraphRAGDataStore(
+            url="bolt://localhost:7687",
+            username="neo4j",
+            password="password",
+            lm_invoker=lm_invoker,  # Optional: Auto-converted to LlamaIndex LLM
+            em_invoker=em_invoker,  # Optional: Auto-converted to LlamaIndex Embedding
+        )
+
+        # Option 2: Use with LlamaIndex LLM/Embeddings directly
+        from llama_index.llms.openai import OpenAI
+        from llama_index.embeddings.openai import OpenAIEmbedding
+
+        store = LlamaIndexNeo4jGraphRAGDataStore(
+            url="bolt://localhost:7687",
+            username="neo4j",
+            password="password",
+        )
+
+        # Perform RAG query
+        results = await store.query("What is the relationship between X and Y?")
+
+        # Delete document data
+        await store.delete_by_document_id("doc123")
+        ```
+    '''
+    neo4j_version_tuple: Incomplete
+    def __init__(self, url: str, username: str, password: str, lm_invoker: BaseLMInvoker | None = None, em_invoker: BaseEMInvoker | None = None, **kwargs: Any) -> None:
+        '''Initialize the LlamaIndexNeo4jGraphRAGDataStore.
+
+        Args:
+            url (str): The Neo4j database URL (e.g., "bolt://localhost:7687").
+            username (str): The Neo4j database username.
+            password (str): The Neo4j database password.
+            lm_invoker (BaseLMInvoker | None, optional): GLLM language model invoker.
+                If provided, it will be automatically converted to a LlamaIndex LLM instance
+                by the parent class. Defaults to None.
+            em_invoker (BaseEMInvoker | None, optional): GLLM embedding model invoker.
+                If provided, it will be automatically converted to a LlamaIndex BaseEmbedding instance
+                by the parent class. Defaults to None.
+            **kwargs (Any): Additional keyword arguments passed to Neo4jPropertyGraphStore.
+        '''
+    async def delete_by_document_id(self, document_id: str, **kwargs: Any) -> None:
+        """Delete nodes and edges by document ID.
+
+        Args:
+            document_id (str): The document ID.
+            **kwargs (Any): Additional keyword arguments.
+        """

gllm_datastore/graph_data_store/nebula_graph_data_store.pyi

@@ -0,0 +1,206 @@
+from _typeshed import Incomplete
+from gllm_datastore.graph_data_store.graph_data_store import BaseGraphDataStore as BaseGraphDataStore
+from typing import Any
+
+class NebulaGraphDataStore(BaseGraphDataStore):
+    '''Implementation of BaseGraphDataStore for Nebula Graph.
+
+    This class provides an interface for graph-based Retrieval-Augmented Generation (RAG)
+    operations on Nebula graph databases.
+
+    Attributes:
+        connection_pool (ConnectionPool): The connection pool for Nebula Graph.
+        space (str): The space name.
+        user (str): The username.
+        password (str): The password.
+        operation_wait_time (int): The timeout in seconds.
+
+    Example:
+        ```python
+        store = NebulaGraphDataStore(
+            url="127.0.0.1",
+            port=9669,
+            user="root",
+            password="nebula",
+            space="testing"
+        )
+        # Perform query
+        results = await store.query("MATCH (n) RETURN n")
+
+        # Create a node
+        node = await store.upsert_node("Person", "name", "John", {"age": 30})
+        ```
+    '''
+    connection_pool: Incomplete
+    space: Incomplete
+    user: Incomplete
+    password: Incomplete
+    operation_wait_time: Incomplete
+    def __init__(self, url: str, port: int, user: str, password: str, space: str, operation_wait_time: int = 5) -> None:
+        """Initialize NebulaGraphDataStore.
+
+        Args:
+            url (str): The URL of the graph store.
+            port (int): The port of the graph store.
+            user (str): The user of the graph store.
+            password (str): The password of the graph store.
+            space (str): The space name.
+            operation_wait_time (int, optional): The operation wait time in seconds. Defaults to 5.
+        """
+    async def upsert_node(self, label: str, identifier_key: str, identifier_value: str, properties: dict[str, Any] | None = None) -> Any:
+        """Upsert a node in the graph.
+
+        Args:
+            label (str): The label of the node.
+            identifier_key (str): The key of the identifier.
+            identifier_value (str): The value of the identifier.
+            properties (dict[str, Any] | None, optional): The properties of the node. Defaults to None.
+
+        Returns:
+            Any: The result of the operation.
+        """
+    async def upsert_relationship(self, node_source_key: str, node_source_value: str, relation: str, node_target_key: str, node_target_value: str, properties: dict[str, Any] | None = None) -> Any:
+        """Upsert a relationship between two nodes in the graph.
+
+        Args:
+            node_source_key (str): The key of the source node.
+            node_source_value (str): The value of the source node.
+            relation (str): The type of the relationship.
+            node_target_key (str): The key of the target node.
+            node_target_value (str): The value of the target node.
+            properties (dict[str, Any] | None, optional): The properties of the relationship. Defaults to None.
+
+        Returns:
+            Any: The result of the operation.
+        """
+    async def delete_node(self, label: str, identifier_key: str, identifier_value: str) -> Any:
+        """Delete a node from the graph.
+
+        Args:
+            label (str): The label of the node.
+            identifier_key (str): The key of the identifier.
+            identifier_value (str): The identifier of the node.
+
+        Returns:
+            Any: The result of the operation.
+        """
+    async def delete_relationship(self, node_source_key: str, node_source_value: str, relation: str, node_target_key: str, node_target_value: str) -> Any:
+        """Delete a relationship between two nodes in the graph.
+
+        Args:
+            node_source_key (str): The key of the source node.
+            node_source_value (str): The identifier of the source node.
+            relation (str): The type of the relationship.
+            node_target_key (str): The key of the target node.
+            node_target_value (str): The identifier of the target node.
+
+        Returns:
+            Any: The result of the operation.
+        """
+    async def query(self, query: str, parameters: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+        """Query the graph store.
+
+        Args:
+            query (str): The query to be executed.
+            parameters (dict[str, Any] | None, optional): The parameters of the query. Defaults to None.
+
+        Returns:
+            list[dict[str, Any]]: The result of the query.
+        """
+    async def traverse_graph(self, node_properties: dict[str, Any], extracted_node_properties: list[str] | None = None, extracted_relationship_properties: list[str] | None = None, depth: int = 3) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+        '''Traverse graph from a node with specified properties, ignoring relationship\'s direction, up to a given depth.
+
+        Example:
+            ```python
+            nodes, relationships = await graph_data_store.traverse_graph(
+                node_properties={"name": "John Doe"},
+                extracted_node_properties=["name", "age"],
+                extracted_relationship_properties=["since"],
+                depth=1
+            )
+            ```
+            Means starting from the node with property `name` equal to "John Doe", traverse
+            the graph up to depth 1, extracting the `name` and `age` properties from nodes
+            and the `since` property from relationships.
+
+            ```python
+            nodes, relationships = await graph_data_store.traverse_graph(
+                node_properties={"name": "John Doe"},
+                depth=2
+            )
+            ```
+            Means starting from the node with property `name` equal to "John Doe", traverse
+            the graph up to depth 2, extracting all properties from nodes and relationships.
+
+        Args:
+            node_properties (dict[str, Any]): The properties of the starting node.
+            extracted_node_properties (list[str] | None, optional): The properties to extract from nodes during
+                traversal. If None or empty list, all node properties will be returned. Defaults to None.
+            extracted_relationship_properties (list[str] | None, optional): The properties to extract from relationships
+                during traversal. If None or empty list, all relationship properties will be returned. Defaults to None.
+            depth (int, optional): The depth of traversal. Defaults to 3.
+
+        Returns:
+            tuple[list[dict[str, Any]], list[dict[str, Any]]]: A tuple containing two lists:
+                - List of nodes with their extracted properties (including the source node).
+                - List of relationships with their extracted properties.
+
+            Example return value:
+            nodes = [
+                {
+                    "id": 1001,
+                    "labels": ["Person"],
+                    "properties": {
+                        "name": "John Doe",
+                        "age": 30,
+                        "occupation": "Engineer"
+                    }
+                },
+                {
+                    "id": 2001,
+                    "labels": ["Company"],
+                    "properties": {
+                        "name": "TechCorp",
+                        "industry": "Technology",
+                        "employees": 500
+                    }
+                }
+            ]
+
+            relationships = [
+                {
+                    "id": 5002,
+                    "type": "FRIEND_OF",
+                    "start_node": 1001,
+                    "end_node": 1002,
+                    "properties": {
+                        "since": "2018-05-20",
+                        "closeness": 8
+                    }
+                }
+            ]
+
+        Raises:
+            ValueError: If node_properties is empty or depth is less than 1.
+        '''
+    async def close(self) -> None:
+        """Close the graph data store."""
+    async def get_nodes(self, label: str | None = None) -> list[dict[str, Any]]:
+        """Get all nodes with optional label filter.
+
+        Args:
+            label (str | None, optional): The label of the nodes. Defaults to None.
+
+        Returns:
+            list[dict[str, Any]]: The result of the query.
+        """
+    async def get_relationships(self, source_value: str | None = None, relation: str | None = None) -> list[dict[str, Any]]:
+        """Get relationships with optional filters.
+
+        Args:
+            source_value (str | None, optional): The source vertex identifier. Defaults to None.
+            relation (str | None, optional): The relationship type. Defaults to None.
+
+        Returns:
+            list[dict[str, Any]]: The result of the query.
+        """