graphmemory 0.1.0__tar.gz

graphmemory-0.1.0/PKG-INFO

@@ -0,0 +1,193 @@
+Metadata-Version: 2.1
+Name: graphmemory
+Version: 0.1.0
+Summary: A package for creating a graph database for use with GraphRAG.
+Home-page: http://pypi.python.org/pypi/graphmemory/
+Author: BradAGI
+Author-email: cavemen_summary_0f@icloud.com
+License: LICENSE.txt
+Keywords: graphrag graph database rag
+Description-Content-Type: text/markdown
+Requires-Dist: duckdb
+Requires-Dist: json
+Requires-Dist: pydantic
+Requires-Dist: os
+Requires-Dist: logging
+
+# GraphMemory - GraphRAG Database
+
+## Overview
+This project provides an embedded graph database implementation with vector similarity search (VSS) using DuckDB. It includes a Python class `GraphRAG` for managing nodes and edges. 
+
+Each node has a unique ID, a JSON data field (any arbitrary dictionary), and a vector of floating point values. 
+
+Each edge has a unique ID, a source node ID, a target node ID, a relationship type, and a weight.
+
+This database can be used for any graph-based RAG application or knowledge graph application.
+
+Vector embeddings can be created using [sentence-transformers](https://www.sbert.net/) or other API based models.
+
+## Requirements
+- Python 3.x
+- DuckDB
+- Pydantic
+
+## Installation
+1. Clone the repository:
+    ```sh
+    git clone https://github.com/bradAGI/GraphMemory
+    ```
+2. Install the required packages:
+    ```sh
+    pip install -r requirements.txt
+    ```
+
+## Usage
+
+### GraphRAG Class
+The `GraphRAG` class provides methods to manage nodes and edges, perform bulk inserts, create indexes, and find nearest neighbors using vector similarity search.
+
+### Auto-Incrementing IDs
+If you do not provide an ID for a node or edge, the database will automatically assign a unique ID.
+
+### Example
+```python
+from graphrag.graphrag import GraphRAG
+from graphrag.models import Node, Edge
+
+
+# Initialize the database from disk (make sure to set vector_length correctly)
+graph_db = GraphRAG(database='graph.db', vector_length=3)
+
+# Insert nodes
+node1 = Node(data={"name": "George Washington", "role": "President"}, vector=[0.1, 0.2, 0.3])
+node1_id = graph_db.insert_node(node1)
+
+node2 = Node(data={"name": "Thomas Jefferson", "role": "Secretary of State"}, vector=[0.4, 0.5, 0.6])
+node2_id = graph_db.insert_node(node2)
+
+# Insert edge
+edge = Edge(source_id=node1_id, target_id=node2_id, relation="served_under", weight=0.5)
+graph_db.insert_edge(edge)
+
+# Print all nodes in the database
+nodes = graph_db.nodes_to_json()
+print("Nodes:", nodes)
+
+# Print all edges in the database
+edges = graph_db.edges_to_json()
+print("Edges:", edges)
+
+# Find connected nodes
+connected_nodes = graph_db.connected_nodes(node1_id)
+print("Connected Nodes:", connected_nodes)
+
+# Find nearest neighbors
+neighbors = graph_db.nearest_neighbors(vector=[0.1, 0.2, 0.3], limit=1)
+print("Nearest Neighbors:", neighbors)
+
+# Insert an edge between the two nodes with a relation
+edge = Edge(source_id=node1_id, target_id=node2_id, relation="served_under", weight=0.5)
+graph_db.insert_edge(edge)
+
+# Define the additional nodes for bulk insert
+nodes = [
+    Node(data={"name": "Alexander Hamilton", "role": "Secretary of the Treasury", "term": "1789–1795"}, vector=[0.7, 0.8, 0.9]),
+    Node(data={"name": "Oliver Wolcott Jr.", "role": "Secretary of the Treasury", "term": "1795–1797"}, vector=[1.6, 1.7, 1.8]),
+]
+
+# Bulk insert nodes
+graph_db.bulk_insert_nodes(nodes)
+
+# Define the additional edges for bulk insert
+edges = [
+    Edge(source_id=nodes[0].id, target_id=nodes[1].id, relation="succeeded_by", weight=0.7),
+    Edge(source_id=nodes[1].id, target_id=nodes[2].id, relation="succeeded_by", weight=0.8)
+]
+
+# Bulk insert edges
+graph_db.bulk_insert_edges(edges)
+
+# Delete a node
+graph_db.delete_node(nodes[-1].id)
+
+# Delete an edge
+graph_db.delete_edge(1, 2)
+
+# Find nearest nodes to a given vector by distance
+neighbors = graph_db.nearest_neighbors(vector=[0.1, 0.2, 0.3], limit=2)
+print("Nearest Neighbors:", neighbors)
+
+# Find connected nodes
+connected_nodes = graph_db.connected_nodes(nodes[1].id)
+print("Connected Nodes:", connected_nodes)
+```
+
+## GraphRAG Class Methods
+
+The `GraphRAG` class provides the following public methods for interacting with the graph database:
+
+1. `__init__(self, database=None, vector_length=3)`
+   - Initializes the database connection and sets up the database schema if necessary.
+
+2. `set_vector_length(self, vector_length)`
+   - Sets the length of the vectors for the nodes in the database.
+
+3. `create_tables(self)`
+   - Creates the necessary database tables for nodes and edges if they do not exist.
+
+4. `insert_node(self, node: Node) -> int`
+   - Inserts a node into the database and returns the node ID.
+
+5. `insert_edge(self, edge: Edge)`
+   - Inserts an edge between two nodes in the database.
+
+6. `bulk_insert_nodes(self, nodes: List[Node]) -> List[Node]`
+   - Performs a bulk insert of multiple nodes into the database.
+
+7. `bulk_insert_edges(self, edges: List[Edge])`
+   - Performs a bulk insert of multiple edges into the database.
+
+8. `delete_node(self, node_id: int)`
+   - Deletes a node and its associated edges from the database.
+
+9. `delete_edge(self, source_id: int, target_id: int)`
+   - Deletes an edge from the database.
+
+10. `create_index(self)`
+    - Creates an index on the node vectors to improve search performance.
+
+11. `nearest_neighbors(self, vector: List[float], limit: int) -> List[Neighbor]`
+    - Finds and returns the nearest neighbor nodes based on vector similarity.
+
+12. `connected_nodes(self, node_id: int) -> List[Node]`
+    - Retrieves all nodes directly connected to the specified node.
+
+13. `nodes_to_json(self)`
+    - Returns a JSON representation of all nodes in the database.
+
+14. `edges_to_json(self)`
+    - Returns a JSON representation of all edges in the database.
+
+15. `get_node(self, node_id: int)`
+    - Retrieves a specific node by its ID.
+
+16. `print_json(self)`
+    - Prints the JSON representation of all nodes and edges in the database.
+
+These methods facilitate the management and querying of the graph database, allowing for efficient data handling and retrieval.
+
+## Testing
+Unit tests are provided in `tests/tests.py`.
+
+### Running Tests
+To run the unit tests, use the following command:
+```sh
+python -m unittest discover -s tests
+```
+
+## License
+This project is licensed under the MIT License. See the LICENSE file for details.
+
+## Contributing
+Contributions are welcome! Please open an issue or submit a pull request.

graphmemory-0.1.0/README.md

@@ -0,0 +1,177 @@
+# GraphMemory - GraphRAG Database
+
+## Overview
+This project provides an embedded graph database implementation with vector similarity search (VSS) using DuckDB. It includes a Python class `GraphRAG` for managing nodes and edges. 
+
+Each node has a unique ID, a JSON data field (any arbitrary dictionary), and a vector of floating point values. 
+
+Each edge has a unique ID, a source node ID, a target node ID, a relationship type, and a weight.
+
+This database can be used for any graph-based RAG application or knowledge graph application.
+
+Vector embeddings can be created using [sentence-transformers](https://www.sbert.net/) or other API based models.
+
+## Requirements
+- Python 3.x
+- DuckDB
+- Pydantic
+
+## Installation
+1. Clone the repository:
+    ```sh
+    git clone https://github.com/bradAGI/GraphMemory
+    ```
+2. Install the required packages:
+    ```sh
+    pip install -r requirements.txt
+    ```
+
+## Usage
+
+### GraphRAG Class
+The `GraphRAG` class provides methods to manage nodes and edges, perform bulk inserts, create indexes, and find nearest neighbors using vector similarity search.
+
+### Auto-Incrementing IDs
+If you do not provide an ID for a node or edge, the database will automatically assign a unique ID.
+
+### Example
+```python
+from graphrag.graphrag import GraphRAG
+from graphrag.models import Node, Edge
+
+
+# Initialize the database from disk (make sure to set vector_length correctly)
+graph_db = GraphRAG(database='graph.db', vector_length=3)
+
+# Insert nodes
+node1 = Node(data={"name": "George Washington", "role": "President"}, vector=[0.1, 0.2, 0.3])
+node1_id = graph_db.insert_node(node1)
+
+node2 = Node(data={"name": "Thomas Jefferson", "role": "Secretary of State"}, vector=[0.4, 0.5, 0.6])
+node2_id = graph_db.insert_node(node2)
+
+# Insert edge
+edge = Edge(source_id=node1_id, target_id=node2_id, relation="served_under", weight=0.5)
+graph_db.insert_edge(edge)
+
+# Print all nodes in the database
+nodes = graph_db.nodes_to_json()
+print("Nodes:", nodes)
+
+# Print all edges in the database
+edges = graph_db.edges_to_json()
+print("Edges:", edges)
+
+# Find connected nodes
+connected_nodes = graph_db.connected_nodes(node1_id)
+print("Connected Nodes:", connected_nodes)
+
+# Find nearest neighbors
+neighbors = graph_db.nearest_neighbors(vector=[0.1, 0.2, 0.3], limit=1)
+print("Nearest Neighbors:", neighbors)
+
+# Insert an edge between the two nodes with a relation
+edge = Edge(source_id=node1_id, target_id=node2_id, relation="served_under", weight=0.5)
+graph_db.insert_edge(edge)
+
+# Define the additional nodes for bulk insert
+nodes = [
+    Node(data={"name": "Alexander Hamilton", "role": "Secretary of the Treasury", "term": "1789–1795"}, vector=[0.7, 0.8, 0.9]),
+    Node(data={"name": "Oliver Wolcott Jr.", "role": "Secretary of the Treasury", "term": "1795–1797"}, vector=[1.6, 1.7, 1.8]),
+]
+
+# Bulk insert nodes
+graph_db.bulk_insert_nodes(nodes)
+
+# Define the additional edges for bulk insert
+edges = [
+    Edge(source_id=nodes[0].id, target_id=nodes[1].id, relation="succeeded_by", weight=0.7),
+    Edge(source_id=nodes[1].id, target_id=nodes[2].id, relation="succeeded_by", weight=0.8)
+]
+
+# Bulk insert edges
+graph_db.bulk_insert_edges(edges)
+
+# Delete a node
+graph_db.delete_node(nodes[-1].id)
+
+# Delete an edge
+graph_db.delete_edge(1, 2)
+
+# Find nearest nodes to a given vector by distance
+neighbors = graph_db.nearest_neighbors(vector=[0.1, 0.2, 0.3], limit=2)
+print("Nearest Neighbors:", neighbors)
+
+# Find connected nodes
+connected_nodes = graph_db.connected_nodes(nodes[1].id)
+print("Connected Nodes:", connected_nodes)
+```
+
+## GraphRAG Class Methods
+
+The `GraphRAG` class provides the following public methods for interacting with the graph database:
+
+1. `__init__(self, database=None, vector_length=3)`
+   - Initializes the database connection and sets up the database schema if necessary.
+
+2. `set_vector_length(self, vector_length)`
+   - Sets the length of the vectors for the nodes in the database.
+
+3. `create_tables(self)`
+   - Creates the necessary database tables for nodes and edges if they do not exist.
+
+4. `insert_node(self, node: Node) -> int`
+   - Inserts a node into the database and returns the node ID.
+
+5. `insert_edge(self, edge: Edge)`
+   - Inserts an edge between two nodes in the database.
+
+6. `bulk_insert_nodes(self, nodes: List[Node]) -> List[Node]`
+   - Performs a bulk insert of multiple nodes into the database.
+
+7. `bulk_insert_edges(self, edges: List[Edge])`
+   - Performs a bulk insert of multiple edges into the database.
+
+8. `delete_node(self, node_id: int)`
+   - Deletes a node and its associated edges from the database.
+
+9. `delete_edge(self, source_id: int, target_id: int)`
+   - Deletes an edge from the database.
+
+10. `create_index(self)`
+    - Creates an index on the node vectors to improve search performance.
+
+11. `nearest_neighbors(self, vector: List[float], limit: int) -> List[Neighbor]`
+    - Finds and returns the nearest neighbor nodes based on vector similarity.
+
+12. `connected_nodes(self, node_id: int) -> List[Node]`
+    - Retrieves all nodes directly connected to the specified node.
+
+13. `nodes_to_json(self)`
+    - Returns a JSON representation of all nodes in the database.
+
+14. `edges_to_json(self)`
+    - Returns a JSON representation of all edges in the database.
+
+15. `get_node(self, node_id: int)`
+    - Retrieves a specific node by its ID.
+
+16. `print_json(self)`
+    - Prints the JSON representation of all nodes and edges in the database.
+
+These methods facilitate the management and querying of the graph database, allowing for efficient data handling and retrieval.
+
+## Testing
+Unit tests are provided in `tests/tests.py`.
+
+### Running Tests
+To run the unit tests, use the following command:
+```sh
+python -m unittest discover -s tests
+```
+
+## License
+This project is licensed under the MIT License. See the LICENSE file for details.
+
+## Contributing
+Contributions are welcome! Please open an issue or submit a pull request.

graphmemory-0.1.0/graphmemory.egg-info/PKG-INFO

@@ -0,0 +1,193 @@
+Metadata-Version: 2.1
+Name: graphmemory
+Version: 0.1.0
+Summary: A package for creating a graph database for use with GraphRAG.
+Home-page: http://pypi.python.org/pypi/graphmemory/
+Author: BradAGI
+Author-email: cavemen_summary_0f@icloud.com
+License: LICENSE.txt
+Keywords: graphrag graph database rag
+Description-Content-Type: text/markdown
+Requires-Dist: duckdb
+Requires-Dist: json
+Requires-Dist: pydantic
+Requires-Dist: os
+Requires-Dist: logging
+
+# GraphMemory - GraphRAG Database
+
+## Overview
+This project provides an embedded graph database implementation with vector similarity search (VSS) using DuckDB. It includes a Python class `GraphRAG` for managing nodes and edges. 
+
+Each node has a unique ID, a JSON data field (any arbitrary dictionary), and a vector of floating point values. 
+
+Each edge has a unique ID, a source node ID, a target node ID, a relationship type, and a weight.
+
+This database can be used for any graph-based RAG application or knowledge graph application.
+
+Vector embeddings can be created using [sentence-transformers](https://www.sbert.net/) or other API based models.
+
+## Requirements
+- Python 3.x
+- DuckDB
+- Pydantic
+
+## Installation
+1. Clone the repository:
+    ```sh
+    git clone https://github.com/bradAGI/GraphMemory
+    ```
+2. Install the required packages:
+    ```sh
+    pip install -r requirements.txt
+    ```
+
+## Usage
+
+### GraphRAG Class
+The `GraphRAG` class provides methods to manage nodes and edges, perform bulk inserts, create indexes, and find nearest neighbors using vector similarity search.
+
+### Auto-Incrementing IDs
+If you do not provide an ID for a node or edge, the database will automatically assign a unique ID.
+
+### Example
+```python
+from graphrag.graphrag import GraphRAG
+from graphrag.models import Node, Edge
+
+
+# Initialize the database from disk (make sure to set vector_length correctly)
+graph_db = GraphRAG(database='graph.db', vector_length=3)
+
+# Insert nodes
+node1 = Node(data={"name": "George Washington", "role": "President"}, vector=[0.1, 0.2, 0.3])
+node1_id = graph_db.insert_node(node1)
+
+node2 = Node(data={"name": "Thomas Jefferson", "role": "Secretary of State"}, vector=[0.4, 0.5, 0.6])
+node2_id = graph_db.insert_node(node2)
+
+# Insert edge
+edge = Edge(source_id=node1_id, target_id=node2_id, relation="served_under", weight=0.5)
+graph_db.insert_edge(edge)
+
+# Print all nodes in the database
+nodes = graph_db.nodes_to_json()
+print("Nodes:", nodes)
+
+# Print all edges in the database
+edges = graph_db.edges_to_json()
+print("Edges:", edges)
+
+# Find connected nodes
+connected_nodes = graph_db.connected_nodes(node1_id)
+print("Connected Nodes:", connected_nodes)
+
+# Find nearest neighbors
+neighbors = graph_db.nearest_neighbors(vector=[0.1, 0.2, 0.3], limit=1)
+print("Nearest Neighbors:", neighbors)
+
+# Insert an edge between the two nodes with a relation
+edge = Edge(source_id=node1_id, target_id=node2_id, relation="served_under", weight=0.5)
+graph_db.insert_edge(edge)
+
+# Define the additional nodes for bulk insert
+nodes = [
+    Node(data={"name": "Alexander Hamilton", "role": "Secretary of the Treasury", "term": "1789–1795"}, vector=[0.7, 0.8, 0.9]),
+    Node(data={"name": "Oliver Wolcott Jr.", "role": "Secretary of the Treasury", "term": "1795–1797"}, vector=[1.6, 1.7, 1.8]),
+]
+
+# Bulk insert nodes
+graph_db.bulk_insert_nodes(nodes)
+
+# Define the additional edges for bulk insert
+edges = [
+    Edge(source_id=nodes[0].id, target_id=nodes[1].id, relation="succeeded_by", weight=0.7),
+    Edge(source_id=nodes[1].id, target_id=nodes[2].id, relation="succeeded_by", weight=0.8)
+]
+
+# Bulk insert edges
+graph_db.bulk_insert_edges(edges)
+
+# Delete a node
+graph_db.delete_node(nodes[-1].id)
+
+# Delete an edge
+graph_db.delete_edge(1, 2)
+
+# Find nearest nodes to a given vector by distance
+neighbors = graph_db.nearest_neighbors(vector=[0.1, 0.2, 0.3], limit=2)
+print("Nearest Neighbors:", neighbors)
+
+# Find connected nodes
+connected_nodes = graph_db.connected_nodes(nodes[1].id)
+print("Connected Nodes:", connected_nodes)
+```
+
+## GraphRAG Class Methods
+
+The `GraphRAG` class provides the following public methods for interacting with the graph database:
+
+1. `__init__(self, database=None, vector_length=3)`
+   - Initializes the database connection and sets up the database schema if necessary.
+
+2. `set_vector_length(self, vector_length)`
+   - Sets the length of the vectors for the nodes in the database.
+
+3. `create_tables(self)`
+   - Creates the necessary database tables for nodes and edges if they do not exist.
+
+4. `insert_node(self, node: Node) -> int`
+   - Inserts a node into the database and returns the node ID.
+
+5. `insert_edge(self, edge: Edge)`
+   - Inserts an edge between two nodes in the database.
+
+6. `bulk_insert_nodes(self, nodes: List[Node]) -> List[Node]`
+   - Performs a bulk insert of multiple nodes into the database.
+
+7. `bulk_insert_edges(self, edges: List[Edge])`
+   - Performs a bulk insert of multiple edges into the database.
+
+8. `delete_node(self, node_id: int)`
+   - Deletes a node and its associated edges from the database.
+
+9. `delete_edge(self, source_id: int, target_id: int)`
+   - Deletes an edge from the database.
+
+10. `create_index(self)`
+    - Creates an index on the node vectors to improve search performance.
+
+11. `nearest_neighbors(self, vector: List[float], limit: int) -> List[Neighbor]`
+    - Finds and returns the nearest neighbor nodes based on vector similarity.
+
+12. `connected_nodes(self, node_id: int) -> List[Node]`
+    - Retrieves all nodes directly connected to the specified node.
+
+13. `nodes_to_json(self)`
+    - Returns a JSON representation of all nodes in the database.
+
+14. `edges_to_json(self)`
+    - Returns a JSON representation of all edges in the database.
+
+15. `get_node(self, node_id: int)`
+    - Retrieves a specific node by its ID.
+
+16. `print_json(self)`
+    - Prints the JSON representation of all nodes and edges in the database.
+
+These methods facilitate the management and querying of the graph database, allowing for efficient data handling and retrieval.
+
+## Testing
+Unit tests are provided in `tests/tests.py`.
+
+### Running Tests
+To run the unit tests, use the following command:
+```sh
+python -m unittest discover -s tests
+```
+
+## License
+This project is licensed under the MIT License. See the LICENSE file for details.
+
+## Contributing
+Contributions are welcome! Please open an issue or submit a pull request.

graphmemory-0.1.0/graphmemory.egg-info/SOURCES.txt

@@ -0,0 +1,11 @@
+README.md
+setup.py
+graphmemory.egg-info/PKG-INFO
+graphmemory.egg-info/SOURCES.txt
+graphmemory.egg-info/dependency_links.txt
+graphmemory.egg-info/requires.txt
+graphmemory.egg-info/top_level.txt
+graphrag/__init__.py
+graphrag/database.py
+graphrag/models.py
+tests/tests.py

graphmemory-0.1.0/graphmemory.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

graphmemory-0.1.0/graphmemory.egg-info/requires.txt

@@ -0,0 +1,5 @@
+duckdb
+json
+pydantic
+os
+logging

graphmemory-0.1.0/graphmemory.egg-info/top_level.txt

@@ -0,0 +1 @@
+graphrag

graphmemory-0.1.0/graphrag/__init__.py

@@ -0,0 +1,3 @@
+from .database import GraphRAG
+from .models import Node, Edge, Neighbor
+

graphmemory-0.1.0/graphrag/database.py

@@ -0,0 +1,258 @@
+import duckdb
+import json
+import os
+import logging
+from contextlib import contextmanager
+from graphrag.models import Node, Edge, Neighbor
+from typing import List
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class GraphRAG:
+    def __init__(self, database=None, vector_length=3):
+        self.database = database
+        self.vector_length = vector_length
+        self.conn = duckdb.connect(database=self.database)
+        self._load_vss_extension()
+        self._configure_database()
+        
+        if database and os.path.exists(database):
+            self.load_database(database)
+        
+        # Check if 'nodes' and 'edges' tables exist, and create them if they do not
+        nodes_exist = self.conn.execute("SELECT 1 FROM information_schema.tables WHERE table_name = 'nodes';").fetchone()
+        edges_exist = self.conn.execute("SELECT 1 FROM information_schema.tables WHERE table_name = 'edges';").fetchone()
+        
+        if not nodes_exist or not edges_exist:
+            self.create_tables()
+            logger.info("Tables created or verified successfully.")
+    
+    def load_database(self, path):
+        if not os.path.exists(path):
+            logger.error(f"Database file not found: {path}")
+            return
+        try:
+            self.conn.execute(f"ATTACH DATABASE '{path}' AS vss;")
+        except duckdb.Error as e:
+            logger.error(f"Error loading database: {e}")
+
+    def _configure_database(self):
+        try:
+            self.conn.execute("SET hnsw_enable_experimental_persistence=true;")
+        except duckdb.Error as e:
+            logger.error(f"Error setting configuration: {e}")
+
+    def _load_vss_extension(self):
+        try:
+            self.conn.execute("INSTALL vss;")
+            self.conn.execute("LOAD vss;")
+        except duckdb.Error as e:
+            logger.error(f"Error loading VSS extension: {e}")
+
+    def set_vector_length(self, vector_length):
+        self.vector_length = vector_length
+        logger.info(f"Vector length set to: {self.vector_length}")
+
+    def create_tables(self):
+        self.conn.execute("CREATE SEQUENCE IF NOT EXISTS seq_id;")
+        self.conn.execute(f"""
+        CREATE TABLE IF NOT EXISTS nodes (
+            id INTEGER DEFAULT nextval('seq_id'),
+            data JSON,
+            vector FLOAT[{self.vector_length}],
+            PRIMARY KEY (id)
+        );
+        """)
+        self.conn.execute("""
+        CREATE TABLE IF NOT EXISTS edges (
+            source_id INTEGER,
+            target_id INTEGER,
+            weight FLOAT,
+            relation TEXT,
+            PRIMARY KEY (source_id, target_id)
+        );
+        """)
+        logger.info("Tables 'nodes' and 'edges' created or already exist.")
+        self.conn.commit()  # Ensure changes are committed
+
+    def insert_node(self, node: Node) -> int:
+        if not self._validate_vector(node.vector):
+            logger.error("Invalid vector: Must be a list of float values.")
+            return -1
+        try:
+            with self.transaction():
+                result = self.conn.execute(
+                    "INSERT INTO nodes (data, vector) VALUES (?, ?) RETURNING id;",
+                    (json.dumps(node.data), node.vector)
+                ).fetchone()
+                if result:
+                    logger.info(f"Node inserted with ID: {result[0]}")
+                return result[0] if result else -1
+        except duckdb.Error as e:
+            logger.error(f"Error during insert node: {e}")
+            return -1
+
+    # Update insert_edge to use Edge model
+    def insert_edge(self, edge: Edge):
+        try:
+            with self.transaction():
+                # Check if source and target nodes exist
+                source_exists = self.conn.execute("SELECT 1 FROM nodes WHERE id = ?", (edge.source_id,)).fetchone()
+                target_exists = self.conn.execute("SELECT 1 FROM nodes WHERE id = ?", (edge.target_id,)).fetchone()
+                if not source_exists or not target_exists:
+                    raise ValueError("Source or target node does not exist.")
+                
+                self.conn.execute("INSERT INTO edges (source_id, target_id, weight, relation) VALUES (?, ?, ?, ?);", (edge.source_id, edge.target_id, edge.weight, edge.relation))
+        except duckdb.Error as e:
+            logger.error(f"Error during insert edge: {e}")
+        except ValueError as e:
+            logger.error(f"Error during insert edge: {e}")
+            raise  # Re-raise the ValueError
+
+    def bulk_insert_nodes(self, nodes: List[Node]) -> List[Node]:
+        try:
+            with self.transaction():
+                for node in nodes:
+                    result = self.conn.execute(
+                        "INSERT INTO nodes (data, vector) VALUES (?, ?) RETURNING id;",
+                        (json.dumps(node.data), node.vector)
+                    ).fetchone()
+                    if result:
+                        node.id = result[0]  # Set the node ID
+                return nodes
+        except duckdb.Error as e:
+            logger.error(f"Error during bulk insert nodes: {e}")
+            return []
+
+    def bulk_insert_edges(self, edges: List[Edge]):
+        try:
+            with self.transaction():
+                self.conn.executemany(
+                    "INSERT INTO edges (source_id, target_id, weight, relation) VALUES (?, ?, ?, ?);",
+                    [(edge.source_id, edge.target_id, edge.weight, edge.relation) for edge in edges]
+                )
+        except duckdb.Error as e:
+            logger.error(f"Error during bulk insert edges: {e}")
+
+    def delete_node(self, node_id: int):
+        try:
+            with self.transaction():
+                self.conn.execute("DELETE FROM nodes WHERE id = ?;", (node_id,))
+                self.conn.execute("DELETE FROM edges WHERE source_id = ? OR target_id = ?;", (node_id, node_id))
+        except duckdb.Error as e:
+            logger.error(f"Error deleting node: {e}")
+
+    def delete_edge(self, source_id: int, target_id: int):
+        try:
+            with self.transaction():
+                self.conn.execute("DELETE FROM edges WHERE source_id = ? AND target_id = ?;", (source_id, target_id))
+        except duckdb.Error as e:
+            logger.error(f"Error deleting edge: {e}")
+
+    def create_index(self):
+        try:
+            self.conn.execute("CREATE INDEX IF NOT EXISTS vss_idx ON nodes USING HNSW(vector);")
+        except duckdb.Error as e:
+            logger.error(f"Error creating index: {e}")
+
+    # Update nearest_neighbors to use vector and limit directly
+    def nearest_neighbors(self, vector: List[float], limit: int) -> List[Neighbor]:
+        if not self._validate_vector(vector):
+            logger.error("Invalid vector: Must be a list of float values.")
+            return []
+
+        query = f"""
+        SELECT id, data, vector, array_distance(vector, CAST(? AS FLOAT[{self.vector_length}])) AS distance
+        FROM nodes
+        ORDER BY distance LIMIT ?;
+        """
+        try:
+            results = self.conn.execute(query, (vector, limit)).fetchall()
+            return [
+                Neighbor(
+                    node=Node(id=row[0], data=json.loads(row[1]), vector=row[2]),
+                    distance=row[3]
+                ) for row in results
+            ]
+        except duckdb.Error as e:
+            logger.error(f"Error fetching nearest neighbors: {e}")
+            return []
+
+    def connected_nodes(self, node_id: int) -> List[Node]:
+        query = """
+        SELECT n.id, n.data, n.vector
+        FROM nodes n
+        WHERE n.id IN (
+                SELECT target_id FROM edges WHERE source_id = CAST(? AS INTEGER)
+                UNION
+                SELECT source_id FROM edges WHERE target_id = CAST(? AS INTEGER)
+            );
+        """
+        try:
+            logger.info(f"Executing query to fetch connected nodes for node_id: {node_id}")
+            results = self.conn.execute(query, (node_id, node_id)).fetchall()
+            if results:
+                connected_nodes = [Node(id=row[0], data=json.loads(row[1]), vector=row[2]) for row in results]
+                logger.info(f"Found {len(connected_nodes)} connected nodes.")
+            else:
+                connected_nodes = []
+                logger.info("No connected nodes found.")
+            return connected_nodes
+        except duckdb.Error as e:
+            logger.error(f"Error fetching connected nodes: {e}")
+            return []
+
+    def nodes_to_json(self):
+        try:
+            nodes = self.conn.execute("SELECT id, data, vector FROM nodes;").fetchall()
+            return [{"id": row[0], "data": json.loads(row[1]), "vector": row[2]} for row in nodes]
+        except duckdb.Error as e:
+            logger.error(f"Error fetching nodes: {e}")
+            return []
+
+    def edges_to_json(self):
+        try:
+            edges = self.conn.execute("SELECT source_id, target_id, weight, relation FROM edges;").fetchall()
+            return [{"source_id": row[0], "target_id": row[1], "weight": row[2], "relation": row[3]} for row in edges]
+        except duckdb.Error as e:
+            logger.error(f"Error fetching edges: {e}")
+            return []
+    
+    # Get node by id
+    def get_node(self, node_id: int):
+        try:
+            node = self.conn.execute("SELECT id, data, vector FROM nodes WHERE id = ?;", (node_id,)).fetchone()
+            return {"id": node[0], "data": json.loads(node[1]), "vector": node[2]}
+        except duckdb.Error as e:
+            logger.error(f"Error fetching node: {e}")
+            return {}
+
+    def print_json(self):
+        nodes_json = self.nodes_to_json()
+        edges_json = self.edges_to_json()
+        print("Nodes JSON:", json.dumps(nodes_json, indent=2))
+        print("Edges JSON:", json.dumps(edges_json, indent=2))
+
+    def _validate_vector(self, vector):
+        return isinstance(vector, list) and len(vector) == self.vector_length and all(isinstance(x, float) for x in vector)
+
+    @contextmanager
+    def transaction(self):
+        try:
+            self.conn.execute("BEGIN TRANSACTION;")
+            yield
+            self.conn.execute("COMMIT;")
+        except Exception as e:
+            self.conn.execute("ROLLBACK;")
+            raise e
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.conn.close()
+        logger.info("Database connection closed.")
+

graphmemory-0.1.0/graphrag/models.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel
+from typing import List, Dict, Any
+
+class Node(BaseModel):
+    id: int = None
+    data: Dict[str, Any]
+    vector: List[float]
+
+class Edge(BaseModel):
+    source_id: int
+    target_id: int
+    relation: str
+    weight: float
+
+class Neighbor(BaseModel):
+    node: Node
+    distance: float

graphmemory-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

graphmemory-0.1.0/setup.py

@@ -0,0 +1,22 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='graphmemory',
+    version='0.1.0',
+    author='BradAGI',
+    author_email='cavemen_summary_0f@icloud.com',
+    packages=find_packages(),
+    url='http://pypi.python.org/pypi/graphmemory/',
+    license='LICENSE.txt',
+    description='A package for creating a graph database for use with GraphRAG.',
+    long_description=open('README.md').read(),
+    long_description_content_type='text/markdown',
+    install_requires=[
+        'duckdb',
+        'json',
+        'pydantic',
+        'os',
+        'logging'
+    ],
+    keywords='graphrag graph database rag'
+)

graphmemory-0.1.0/tests/tests.py

@@ -0,0 +1,239 @@
+import sys
+import os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))
+import duckdb
+import unittest
+from src.graphrag import GraphRAG
+from src.models import Node, Edge, Neighbor
+from pydantic import ValidationError
+
+class TestGraphRAG(unittest.TestCase):
+
+    def setUp(self):
+        self.db = GraphRAG(database=':memory:', vector_length=3)
+
+    def test_insert_node(self):
+        node = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        node_id = self.db.insert_node(node)
+        result = self.db.conn.execute("SELECT * FROM nodes WHERE id = ?", (node_id,)).fetchone()
+        self.assertIsNotNone(result)
+        self.assertEqual(result[0], node_id)
+        self.assertEqual(result[1], '{"name": "node1"}')
+        self.assertAlmostEqual(result[2][0], 0.1, places=7)
+        self.assertAlmostEqual(result[2][1], 0.2, places=7)
+        self.assertAlmostEqual(result[2][2], 0.3, places=7)
+
+    def test_insert_invalid_node(self):
+        with self.assertRaises(ValidationError):
+            node = Node(data={"name": "node1"}, vector=[0.1, 0.2, "invalid"])
+            self.db.insert_node(node)
+
+    def test_insert_node_invalid_vector_length(self):
+        node = Node(data={"name": "node1"}, vector=[0.1, 0.2])
+        result = self.db.insert_node(node)
+        self.assertEqual(result, -1)
+
+    def test_insert_edge(self):
+        node1 = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        node2 = Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6])
+        node1_id = self.db.insert_node(node1)
+        node2_id = self.db.insert_node(node2)
+        edge = Edge(source_id=node1_id, target_id=node2_id, weight=0.5, relation="friendship")
+        self.db.insert_edge(edge)
+        result = self.db.conn.execute("SELECT source_id, target_id, weight, relation FROM edges WHERE source_id = ? AND target_id = ?", (node1_id, node2_id)).fetchone()
+        self.assertIsNotNone(result)
+        self.assertEqual(result[0], node1_id)
+        self.assertEqual(result[1], node2_id)
+        self.assertEqual(result[2], 0.5)
+        self.assertEqual(result[3], "friendship")
+
+    def test_insert_edge_non_existent_nodes(self):
+        edge = Edge(source_id=999, target_id=1000, weight=0.5, relation="friendship")
+        with self.assertRaises(ValueError):
+            self.db.insert_edge(edge)
+
+    def test_nearest_neighbors_empty_db(self):
+        neighbors = self.db.nearest_neighbors(vector=[0.1, 0.2, 0.3], limit=2)
+        self.assertEqual(len(neighbors), 0)
+
+    def test_nearest_neighbors(self):
+        nodes = [
+            Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3]),
+            Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6]),
+            Node(data={"name": "node3"}, vector=[0.7, 0.8, 0.9])
+        ]
+        for node in nodes:
+            self.db.insert_node(node)
+        self.db.create_index()
+        neighbors = self.db.nearest_neighbors(vector=[0.1, 0.2, 0.3], limit=2)
+        self.assertEqual(len(neighbors), 2)
+        self.assertEqual(neighbors[0].node.id, 1)
+        self.assertEqual(neighbors[1].node.id, 2)
+
+    def test_nodes_to_json(self):
+        node = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        self.db.insert_node(node)
+        nodes_json = self.db.nodes_to_json()
+        self.assertEqual(len(nodes_json), 1)
+        self.assertEqual(nodes_json[0]['id'], 1)
+        self.assertEqual(nodes_json[0]['data'], {"name": "node1"})
+        self.assertAlmostEqual(nodes_json[0]['vector'][0], 0.1, places=7)
+        self.assertAlmostEqual(nodes_json[0]['vector'][1], 0.2, places=7)
+        self.assertAlmostEqual(nodes_json[0]['vector'][2], 0.3, places=7)
+
+    def test_edges_to_json(self):
+        node1 = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        node2 = Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6])
+        self.db.insert_node(node1)
+        self.db.insert_node(node2)
+        edge = Edge(source_id=1, target_id=2, weight=0.5, relation="friendship")
+        self.db.insert_edge(edge)
+        edges_json = self.db.edges_to_json()
+        self.assertEqual(len(edges_json), 1)
+        self.assertEqual(edges_json[0]['source_id'], 1)
+        self.assertEqual(edges_json[0]['target_id'], 2)
+        self.assertEqual(edges_json[0]['weight'], 0.5)
+        self.assertEqual(edges_json[0]['relation'], "friendship")
+
+    def test_bulk_insert_nodes(self):
+        nodes = [
+            Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3]),
+            Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6])
+        ]
+        inserted_nodes = self.db.bulk_insert_nodes(nodes)
+        self.assertEqual(len(inserted_nodes), 2)
+        self.assertTrue(all(node.id is not None for node in inserted_nodes))
+
+    def test_bulk_insert_edges(self):
+        node1 = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        node2 = Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6])
+        node3 = Node(data={"name": "node3"}, vector=[0.7, 0.8, 0.9])
+        self.db.bulk_insert_nodes([node1, node2, node3])
+        edges = [
+            Edge(source_id=1, target_id=2, weight=0.5, relation="friendship"),
+            Edge(source_id=2, target_id=3, weight=0.7, relation="colleague")
+        ]
+        self.db.bulk_insert_edges(edges)
+        result = self.db.conn.execute("SELECT * FROM edges").fetchall()
+        self.assertEqual(len(result), 2)
+
+    def test_delete_node(self):
+        node = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        node_id = self.db.insert_node(node)
+        self.db.delete_node(node_id)
+        result = self.db.conn.execute("SELECT * FROM nodes WHERE id = ?", (node_id,)).fetchone()
+        self.assertIsNone(result)
+
+    def test_delete_non_existent_node(self):
+        self.db.delete_node(999)
+        result = self.db.conn.execute("SELECT * FROM nodes WHERE id = 999").fetchone()
+        self.assertIsNone(result)
+
+    def test_delete_edge(self):
+        node1 = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        node2 = Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6])
+        self.db.insert_node(node1)
+        self.db.insert_node(node2)
+        edge = Edge(source_id=1, target_id=2, weight=0.5, relation="friendship")
+        self.db.insert_edge(edge)
+        self.db.delete_edge(1, 2)
+        result = self.db.conn.execute("SELECT * FROM edges WHERE source_id = 1 AND target_id = 2").fetchone()
+        self.assertIsNone(result)
+
+    def test_delete_non_existent_edge(self):
+        self.db.delete_edge(999, 1000)
+        result = self.db.conn.execute("SELECT * FROM edges WHERE source_id = 999 AND target_id = 1000").fetchone()
+        self.assertIsNone(result)
+
+    def test_transaction_handling(self):
+        try:
+            with self.db.transaction():
+                self.db.conn.execute("INSERT INTO nodes (data, vector) VALUES ('{\"name\": \"node1\"}', [0.1, 0.2, 0.3]);")
+                raise Exception("Force rollback")
+        except:
+            pass
+        result = self.db.conn.execute("SELECT * FROM nodes WHERE data = '{\"name\": \"node1\"}'").fetchone()
+        self.assertIsNone(result)
+
+    def test_auto_increment_node_id(self):
+        node1 = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        node2 = Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6])
+        id1 = self.db.insert_node(node1)
+        id2 = self.db.insert_node(node2)
+        
+        self.assertEqual(id1, 1)
+        self.assertEqual(id2, 2)
+
+    def test_auto_increment_edge_id(self):
+        node1 = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        node2 = Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6])
+        id1 = self.db.insert_node(node1)
+        id2 = self.db.insert_node(node2)
+        
+        edge1 = Edge(source_id=id1, target_id=id2, weight=0.5, relation="friendship")
+        edge2 = Edge(source_id=id2, target_id=id1, weight=0.7, relation="colleague")
+        self.db.insert_edge(edge1)
+        self.db.insert_edge(edge2)
+        
+        result1 = self.db.conn.execute("SELECT source_id, target_id FROM edges WHERE source_id = ? AND target_id = ?", (id1, id2)).fetchone()
+        result2 = self.db.conn.execute("SELECT source_id, target_id FROM edges WHERE source_id = ? AND target_id = ?", (id2, id1)).fetchone()
+        
+        self.assertIsNotNone(result1)
+        self.assertIsNotNone(result2)
+        self.assertEqual(result1[0], id1)
+        self.assertEqual(result1[1], id2)
+        self.assertEqual(result2[0], id2)
+        self.assertEqual(result2[1], id1)
+
+    def tearDown(self):
+        self.db.conn.close()
+        # Ensure any test directories are cleaned up
+        import shutil
+        if os.path.exists('test_db_parquet'):
+            shutil.rmtree('test_db_parquet')
+        if os.path.exists('test_db_csv'):
+            shutil.rmtree('test_db_csv')
+
+class TestGraphRAGGetConnectedNodes(unittest.TestCase):
+
+    def setUp(self):
+        # Initialize the GraphRAG instance with an in-memory database for testing
+        self.db = GraphRAG(database=':memory:', vector_length=3)
+        self.db.create_tables()  # Ensure tables are created
+
+        # Insert nodes and edges for testing
+        self.node1 = Node(data={"name": "node1"}, vector=[0.1, 0.2, 0.3])
+        self.node2 = Node(data={"name": "node2"}, vector=[0.4, 0.5, 0.6])
+        self.node3 = Node(data={"name": "node3"}, vector=[0.7, 0.8, 0.9])
+        self.node1_id = self.db.insert_node(self.node1)
+        self.node2_id = self.db.insert_node(self.node2)
+        self.node3_id = self.db.insert_node(self.node3)
+
+        # Insert edges
+        self.db.insert_edge(Edge(source_id=self.node1_id, target_id=self.node2_id, weight=0.5, relation="friendship"))
+        self.db.insert_edge(Edge(source_id=self.node2_id, target_id=self.node3_id, weight=0.7, relation="colleague"))
+
+    def test_get_connected_nodes(self):
+        # Test for node1, should have node2 as connected
+        connected_nodes = self.db.get_connected_nodes(self.node1_id)
+        self.assertEqual(len(connected_nodes), 1)
+        self.assertEqual(connected_nodes[0].id, self.node2_id)
+
+        # Test for node2, should have node1 and node3 as connected
+        connected_nodes = self.db.get_connected_nodes(self.node2_id)
+        self.assertEqual(len(connected_nodes), 2)
+        connected_ids = [node.id for node in connected_nodes]
+        self.assertIn(self.node1_id, connected_ids)
+        self.assertIn(self.node3_id, connected_ids)
+
+        # Test for node3, should have node2 as connected
+        connected_nodes = self.db.get_connected_nodes(self.node3_id)
+        self.assertEqual(len(connected_nodes), 1)
+        self.assertEqual(connected_nodes[0].id, self.node2_id)
+
+    def tearDown(self):
+        self.db.conn.close()
+
+if __name__ == '__main__':
+    unittest.main()
+