graphrag_sdk 0.1.0__tar.gz

graphrag_sdk-0.1.0/PKG-INFO

@@ -0,0 +1,155 @@
+Metadata-Version: 2.1
+Name: graphrag_sdk
+Version: 0.1.0
+Summary: 
+Home-page: http://RAG-SDK.readthedocs.io
+License: MIT
+Keywords: RAG,Graph-RAG,KnowledgeGraph
+Author: FalkorDB
+Author-email: info@falkordb.com
+Requires-Python: >=3.12,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: bs4 (>=0.0.2,<0.0.3)
+Requires-Dist: falkordb (>=1.0.4,<2.0.0)
+Requires-Dist: openai (>=1.30.3,<2.0.0)
+Requires-Dist: sqlalchemy (>=2.0.30,<3.0.0)
+Requires-Dist: typing-extensions (>=4.12.1,<5.0.0)
+Project-URL: Repository, https://github.com/FalkorDB/GraphRAG-SDK
+Description-Content-Type: text/markdown
+
+# GraphRAG-SDK
+
+## Example
+
+### Prerequisites
+GraphRAG-SDK relies on [FalkorDB](http://falkordb.com) as its graph engine and works with OpenAI.
+
+Start FalkorDB locally:
+```sh
+docker run -p 6379:6379 -it --rm -v ./data:/data falkordb/falkordb:edge
+```
+
+Export your OpenAI API KEY:
+```sh
+export OPENAI_API_KEY=<YOUR_OPENAI_KEY>
+```
+
+```python
+from graphrag_sdk.schema import Schema
+from graphrag_sdk import KnowledgeGraph, Source
+
+# Auto generate graph schema from unstructured data
+sources = [Source("./data/the_matrix.txt")]
+s = Schema.auto_detect(sources)
+
+# Create a knowledge graph based on schema
+g = KnowledgeGraph("IMDB", schema=s)
+g.process_sources(sources)
+
+# Query your data
+question = "Name a few actors who've played in 'The Matrix'"
+answer, messages = g.ask(question)
+print(f"Answer: {answer}")
+
+# Output:
+# Answer: A few actors who've played in 'The Matrix' are:
+# - Keanu Reeves
+# - Laurence Fishburne
+# - Carrie-Anne Moss
+# - Hugo Weaving
+```
+
+## Introduction
+GraphRAG-SDK provides easy-to-use tooling to get you up and running with your own
+Graph-RAG solution.
+
+There are two main components:
+### Schema
+A `schema` represents the types of entities and relationships within your data.
+For example, the main entities in your data are:  Movies, Actors, and Directors.
+These are interconnected via `ACT` and `DIRECTED` edges.
+
+Two approaches to schema creation are available:
+
+#### Manual schema creation
+Use this method when you know exactly how your data should be structured.
+
+```python
+s = Schema()
+s.add_entity('Actor').add_attribute('name', str, unique=True)
+s.add_entity('Movie').add_attribute('title', str, unique=True)
+s.add_relation("ACTED", 'Actor', 'Movie')
+
+print(f"Schema: {s.to_JSON()}")
+
+# Output:
+# Schema: {"entities": [
+#   {"name": "Actor",
+#    "attributes": [
+#       {"name": "name", "type": "str", "desc": "Actor's name", "unique": true, "mandatory": false}]},
+#   {"name": "Movie",
+#    "attributes": [
+#       {"name": "title", "type": "str", "desc": "Movie's title", "unique": true, "mandatory": false}]}],
+#   "relations": [{"name": "ACTED", "src": "Actor", "dest": "Movie"}]}
+```
+
+#### Automatic schema creation
+Use this method to discover the main entities and relationships within your data.
+Once the schema is discovered, you can adjust it to your liking.
+
+```python
+sources = [Source("./data/madoff.txt")]
+s = Schema.auto_detect(sources)
+json_schema = s.to_JSON()
+print(f"Schema: {json_schema}")
+
+# Adjust json to your liking and reload the schema from JSON
+json_schema['entities'].append({"name": "Movie", "attributes": [ {"name": "Title", "type": "str", "desc": "Movie's title", "unique": true, "mandatory": true}])
+
+# Recreate refined schema
+s = Schema.from_JSON(json_schema)
+
+# Output:
+# Schema: {"entities": [
+#   {"name": "Actor", "attributes": [
+#       {"name": "Name", "type": "str", "desc": "Actor's Name", "unique": true, "mandatory": true}]},
+#   {"name": "Critic", "attributes": [
+#       {"name": "TopCritic", "type": "bool", "desc": "Critic's TopCritic", "unique": false, "mandatory": false},
+#       {"name": "Name", "type": "str", "desc": "Critic's Name", "unique": true, "mandatory": true},
+#       {"name": "Publication", "type": "str", "desc": "Critic's Publication", "unique": false, "mandatory": false}]}...
+```
+
+### KnowledgeGraph
+A `KnowledgeGraph` holds the actual entities and relationships within your data.
+Once constructed, it serves as the backbone of your RAG solution. A Large Language Model will query your knowledge-graph to build a precise context on which its answer will be based.
+
+```python
+# Create a knowledge graph based on a schema.
+g = KnowledgeGraph("IMDB", host="127.0.0.1", port=6379, schema=s)
+g.process_sources([Source("./data/the_matrix.txt"), Source("./data/wework.txt")])
+
+# Query your data
+question = "Name a few actors who've played in 'The Matrix'"
+answer, messages = g.ask(question)
+print(f"Answer: {answer}")
+
+question = "List additional actors"
+answer, messages = g.ask(question, messages)
+print(f"Answer: {answer}")
+
+# Outputs:
+# Answer: A few actors who've played in 'The Matrix' are:
+# - Keanu Reeves
+# - Laurence Fishburne
+# - Carrie-Anne Moss
+# - Hugo Weaving
+#
+# Answer: A few additional actors who've played in 'The Matrix' are:
+# - Joe Pantoliano
+# - Marcus Chong
+# - Paul Goddard
+# - Gloria Foster
+```
+

graphrag_sdk-0.1.0/README.md

@@ -0,0 +1,133 @@
+# GraphRAG-SDK
+
+## Example
+
+### Prerequisites
+GraphRAG-SDK relies on [FalkorDB](http://falkordb.com) as its graph engine and works with OpenAI.
+
+Start FalkorDB locally:
+```sh
+docker run -p 6379:6379 -it --rm -v ./data:/data falkordb/falkordb:edge
+```
+
+Export your OpenAI API KEY:
+```sh
+export OPENAI_API_KEY=<YOUR_OPENAI_KEY>
+```
+
+```python
+from graphrag_sdk.schema import Schema
+from graphrag_sdk import KnowledgeGraph, Source
+
+# Auto generate graph schema from unstructured data
+sources = [Source("./data/the_matrix.txt")]
+s = Schema.auto_detect(sources)
+
+# Create a knowledge graph based on schema
+g = KnowledgeGraph("IMDB", schema=s)
+g.process_sources(sources)
+
+# Query your data
+question = "Name a few actors who've played in 'The Matrix'"
+answer, messages = g.ask(question)
+print(f"Answer: {answer}")
+
+# Output:
+# Answer: A few actors who've played in 'The Matrix' are:
+# - Keanu Reeves
+# - Laurence Fishburne
+# - Carrie-Anne Moss
+# - Hugo Weaving
+```
+
+## Introduction
+GraphRAG-SDK provides easy-to-use tooling to get you up and running with your own
+Graph-RAG solution.
+
+There are two main components:
+### Schema
+A `schema` represents the types of entities and relationships within your data.
+For example, the main entities in your data are:  Movies, Actors, and Directors.
+These are interconnected via `ACT` and `DIRECTED` edges.
+
+Two approaches to schema creation are available:
+
+#### Manual schema creation
+Use this method when you know exactly how your data should be structured.
+
+```python
+s = Schema()
+s.add_entity('Actor').add_attribute('name', str, unique=True)
+s.add_entity('Movie').add_attribute('title', str, unique=True)
+s.add_relation("ACTED", 'Actor', 'Movie')
+
+print(f"Schema: {s.to_JSON()}")
+
+# Output:
+# Schema: {"entities": [
+#   {"name": "Actor",
+#    "attributes": [
+#       {"name": "name", "type": "str", "desc": "Actor's name", "unique": true, "mandatory": false}]},
+#   {"name": "Movie",
+#    "attributes": [
+#       {"name": "title", "type": "str", "desc": "Movie's title", "unique": true, "mandatory": false}]}],
+#   "relations": [{"name": "ACTED", "src": "Actor", "dest": "Movie"}]}
+```
+
+#### Automatic schema creation
+Use this method to discover the main entities and relationships within your data.
+Once the schema is discovered, you can adjust it to your liking.
+
+```python
+sources = [Source("./data/madoff.txt")]
+s = Schema.auto_detect(sources)
+json_schema = s.to_JSON()
+print(f"Schema: {json_schema}")
+
+# Adjust json to your liking and reload the schema from JSON
+json_schema['entities'].append({"name": "Movie", "attributes": [ {"name": "Title", "type": "str", "desc": "Movie's title", "unique": true, "mandatory": true}])
+
+# Recreate refined schema
+s = Schema.from_JSON(json_schema)
+
+# Output:
+# Schema: {"entities": [
+#   {"name": "Actor", "attributes": [
+#       {"name": "Name", "type": "str", "desc": "Actor's Name", "unique": true, "mandatory": true}]},
+#   {"name": "Critic", "attributes": [
+#       {"name": "TopCritic", "type": "bool", "desc": "Critic's TopCritic", "unique": false, "mandatory": false},
+#       {"name": "Name", "type": "str", "desc": "Critic's Name", "unique": true, "mandatory": true},
+#       {"name": "Publication", "type": "str", "desc": "Critic's Publication", "unique": false, "mandatory": false}]}...
+```
+
+### KnowledgeGraph
+A `KnowledgeGraph` holds the actual entities and relationships within your data.
+Once constructed, it serves as the backbone of your RAG solution. A Large Language Model will query your knowledge-graph to build a precise context on which its answer will be based.
+
+```python
+# Create a knowledge graph based on a schema.
+g = KnowledgeGraph("IMDB", host="127.0.0.1", port=6379, schema=s)
+g.process_sources([Source("./data/the_matrix.txt"), Source("./data/wework.txt")])
+
+# Query your data
+question = "Name a few actors who've played in 'The Matrix'"
+answer, messages = g.ask(question)
+print(f"Answer: {answer}")
+
+question = "List additional actors"
+answer, messages = g.ask(question, messages)
+print(f"Answer: {answer}")
+
+# Outputs:
+# Answer: A few actors who've played in 'The Matrix' are:
+# - Keanu Reeves
+# - Laurence Fishburne
+# - Carrie-Anne Moss
+# - Hugo Weaving
+#
+# Answer: A few additional actors who've played in 'The Matrix' are:
+# - Joe Pantoliano
+# - Marcus Chong
+# - Paul Goddard
+# - Gloria Foster
+```

graphrag_sdk-0.1.0/graphrag_sdk/__init__.py

@@ -0,0 +1,2 @@
+from .kg import KnowledgeGraph
+from .source import Source

graphrag_sdk-0.1.0/graphrag_sdk/kg.py

@@ -0,0 +1,155 @@
+import os
+import queue
+import concurrent.futures
+from typing_extensions import AbstractSet
+
+from .prompts import *
+from .template import TEMPLATE
+from openai import OpenAI
+from falkordb import FalkorDB
+from .query_graph import query
+from .schema.schema import Schema
+from .source import Source, AbstractSource
+from .schema.tools_gen import schema_to_tools
+from .schema.functions_gen import schema_to_functions
+
+class KnowledgeGraph(object):
+    """Knowledge Graph model data as a network of entities and relations
+    To create one it is best to provide a schema which will define the graph's ontology
+    In addition to a set of sources from which entities and relations will be extracted.
+    """
+
+    def __init__(self, name:str, host:str="127.0.0.1", port:int=6379, username:str|None=None, password:str|None=None,
+        model:str="gpt-4-1106-preview", schema:Schema|None=None):
+        """
+        Initialize Knowledge Graph
+
+        Parameters:
+            name (str): Knowledge graph name.
+            host (str): FalkorDB hostname.
+            port (int): FalkorDB port number.
+            username (str|None): FalkorDB username.
+            password (str|None): FalkorDB password.
+            model (str): OpenAI model to use.
+        """
+
+        if not isinstance(name, str) or name == "":
+            raise Exception("name should be a non empty string")
+
+        # connect to database
+        self.db = FalkorDB(host=host, port=port, username=username, password=password)
+        self.graph = self.db.select_graph(name)
+
+        self.name = name
+        self.model = model
+        self.schema = schema
+        self.client = None # OpenAI client
+        self.sources = set([])
+        self.ontolegy_ratio = 0.1 # Sampaling ratio for ontolegy detection
+
+        # in case schema is None
+        # try to load schema from FalkorDB
+        if schema is None:
+            schema_name = self._schema_name()
+            if schema_name in self.db.list_graphs():
+                schema_graph = self.db.select_graph()
+                self.schema = self.schema.from_graph(schema_graph)
+
+    def list_sources(self) -> list[AbstractSource]:
+        """
+        List of sources associated with knowledge graph
+
+        Returns:
+            list[AbstractSource]: sources
+        """
+
+        return [s.source for s in self.sources]
+
+    def process_sources(self, sources: list[AbstractSource]) -> None:
+        """
+        Add entities and relations found in sources into the knowledge-graph
+
+        Parameters:
+            sources (list[AbstractSource]): list of sources to extract knowledge from
+        """
+
+        # Make sure knowledge graph is created.
+        if not 'build_graph_from_sources' in globals() or self.client is None:
+            self._create()
+
+        # Run assistant
+        result = build_graph_from_sources(self, self.client, sources)
+
+        # Add processed sources
+        for src in sources:
+            self.sources.add(src)
+
+    def _schema_name(self) -> str:
+        """
+        Generate a name for the schema based on the Knowledge Graph name
+        """
+
+        return f"{self.name}_schema"
+
+    def _createKGWithSchema(self):
+        # Save schema as an ontology graph
+        schema_graph = self.db.select_graph(self._schema_name())
+        self.schema.save_graph(schema_graph)
+
+        # Create Knowledge Graph construction assistant
+
+        # Generate assistant's toolings
+        tools = schema_to_tools(self.schema)
+
+        # Generate functions from schema
+        funcs = schema_to_functions(self.schema)
+        functions_names = [func[0] for func in funcs]
+        functions_impl  = [func[1] for func in funcs]
+
+        # Generate assistant
+        # Generate code from template
+        code = TEMPLATE
+        code = code.replace("<TOOLS>", ",\n".join([str(tool) for tool in tools]))
+        code = code.replace("<GRAPH_ID>", "\"" + self.name + "\"")
+        code = code.replace("<FUNCTIONS>", "\n".join(functions_impl))
+        code = code.replace("<AVAILABLE_FUNCTIONS>", '{' + ",\n".join([f"'{f}': {f}" for f in functions_names]) +'}')
+
+        # Execute the code string in the global namespace
+        # Introduces the 'build_graph_from_sources' function
+        # TODO: note multiple calls to _createKGWithSchema will overide one another.
+        exec(code, globals())
+
+    def _createKGWithoutSchema(self):
+        raise Exception("Can not create knowledge graph, schema missing")
+
+    def _create(self) -> None:
+        """
+        Create knowledge graph
+        Extract entities and relations from sources
+        """
+        self.client = OpenAI()
+
+        if self.schema is not None:
+            self._createKGWithSchema()
+        else:
+            self._createKGWithoutSchema()
+
+    def ask(self, question:str, history:list|None=None) -> tuple[str, list]:
+        """
+        Query the knowledge graph using natural language
+        if the query is asked as part of a longer conversation make sure to
+        include past history.
+
+        Rreturns:
+            tuple[str, list]: answer, conversation history
+
+         Example:
+            >>> ans, history = kg.ask("Which actor has the most oscars")
+            >>> ans, history = kg.ask("List a few movies in which that actored played in", history)
+        """
+
+        if history is None:
+            history = []
+
+        answer = query(self, question, history, self.model)
+        return (answer, history)

graphrag_sdk-0.1.0/graphrag_sdk/prompts.py

@@ -0,0 +1,76 @@
+ONTOLOGY_DETECTION_SYSTEM_PROMPT = """Create SQL statement which will capture the different entities and the relations between them in the provided text.
+Each table should represent either a single entity type with the entity's attributes as the table columns or a relation between two entities, using foreign keys.
+For example an IMDB document is likely to produce 5 tables: Movie, Actor and Director each with its own columns and ACTED and DIRECTED tables.
+Do not create an ID primary key column but set the most compelling attribute as the entity's primary key
+Make sure to call the run_sql function provided in your tools set for each sql statement generated."""
+
+KNOWLEDGE_EXTRACTION_SYSTEM_PROMPT = """You are a helpful assistant with the goal of extracting
+entities and relationships from text"""
+
+SQL_FOLDING_PROMPT = """The following two CREATE TABLE SQL statments are
+creating the same table name, consolidate the two into a single
+CREATE TABLE SQL statment.
+Remove all duplicated columns.
+Make sure the number of PRIMARY KEYs and FOREIGN KEYs is maintained.
+Reply only with the result SQL statment, the reply should start with:
+CREATE TABLE and end with );
+For example: given the two CREATE TABLE SQL:
+CREATE TABLE Person (FirstName string, LastName string, Age int);
+CREATE TABLE Person (SurName string, ForeName string, Height int);
+The response should be:
+CREATE TABLE Person (FirstName string, LastName string, Age int, Height int);
+"""
+
+SQL_FOCUS_PROMPT = """Given a CREATE TABLE SQL statment the table name
+ represents some entity, in case the number of columns in the table is
+ greater than 5, remove columns which are less descriptive of the entity.
+ Do NOT remove columns which are either a PRIMARY KEY or a FOREIGN KEY.
+ Make sure to leave at least 5 columns.
+ Do NOT rename the table name.
+ Reply only with the result SQL statment, the reply should start with:
+ CREATE TABLE and end with );
+ For example: given the CREATE TABLE SQL:
+ CREATE TABLE Movie (
+     Title string PRIMARY KEY,
+     Genre string,
+     ReleaseYear int,
+     Rating string,
+     Language string,
+     BoxOffice string,
+     streaming_date STRING,
+     runtime STRING
+ );
+ The response should be:
+ CREATE TABLE Movie (
+     Title string PRIMARY KEY,
+     Genre string,
+     ReleaseYear int,
+     Rating string,
+     runtime STRING
+ );
+ Another example: given the CREATE TABLE SQL:
+     CREATE TABLE Episode (
+         show_title VARCHAR(255),
+         show_year VARCHAR(4),
+         season INT,
+         episode_number INT,
+         episode_title VARCHAR(255),
+         air_date DATE,
+         episode_rating DECIMAL(2, 1),
+         FOREIGN KEY (show_title, show_year) REFERENCES TVShow(title, year)
+     );
+ You are not allowed to remove either show_title or show_year as these are part of the table's
+ foreign key.
+ """
+
+RELATIONSHIP_NAMING_PROMPT = """"You are a Cypher expert, help suggest names for
+relationship types connecting source node to destination node.
+Respond only with the suggested name.
+For example given source node of type 'Actor' and destination node of type 'Movie'
+Replay with: ACTED"""
+
+FIND_UNIQUE_ATTRIBUTE_PROMPT = """From the list of attributes describing a graph node
+pick the most suitble attribute to act as a unique identifier of the node.
+Respond only with the suggested attribute name
+For example given a list of attributes: series_title, air_date, season_number, episode_number, title
+Describing the graph entity "Episode" Reply with: title"""

graphrag_sdk-0.1.0/graphrag_sdk/query_graph.py

@@ -0,0 +1,144 @@
+import json
+from openai import OpenAI
+
+# Equipet assistant with tools:
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "run_cypher_query",
+            "description": "Run a Cypher query against the graph",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "query to execute",
+                    },
+                },
+                "required": ["query"],
+            },
+        },
+    }
+]
+
+def run_cypher_query(g, q):
+    return g.query(q).result_set
+
+# convert a schema into a string representation
+def _graph_schema_to_prompt(s) -> str:
+    # Build a string description of graph schema
+    desc = "The knowladge graph schema is as follows:\n"
+
+    #--------------------------------------------------------------------------
+    # Describe labels
+    #--------------------------------------------------------------------------
+
+    # list labels
+    desc = desc + "The graph contains the following node labels:\n";
+    for e in s.entities:
+        desc += e + "\n"
+
+    # specify attributes associated with each entity
+    for e in s.entities:
+        e = s.entities[e]
+
+        lbl_name   = e.name
+        attributes = e.attributes
+        attr_count = len(attributes)
+
+        if (attr_count == 0):
+            desc += f"The {lbl_name} label has no attributes\n"
+        else:
+            desc += f"The {lbl_name} label is associated with the following attribute(s):\n"
+            for attr in attributes:
+                attr_name = attr.name
+                attr_type = attr.type.__name__
+                desc += f"{attr_name} which is of type {attr_type}\n"
+
+    desc += "The graph contains the following relationship types:\n"
+
+    #--------------------------------------------------------------------------
+    # Describe relationships
+    #--------------------------------------------------------------------------
+
+    # list relations
+    for rel in s.relations:
+        rel = s.relations[rel]
+
+        src               = rel.src.name
+        relationship_type = rel.name
+        dest              = rel.dest.name
+        desc += f"{relationship_type}\n"
+        desc += f"The {relationship_type} relationship connects the following labels:\n"
+        desc += f"{src} is connected via {relationship_type} to {dest}\n"
+
+    desc += "This is the end of the knowladge graph schema description.\n"
+
+    return desc;
+
+def run_conversation(g, question, messages, model) -> str:
+    client = OpenAI()
+
+    # Step 1: send the conversation and available functions to the model
+    messages.append({"role": "user", "content": question})
+
+    response = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        tools=tools,
+        tool_choice="auto",  # auto is default, but we'll be explicit
+    )
+
+    response_message = response.choices[0].message
+    tool_calls = response_message.tool_calls
+
+    # Step 2: check if the model wanted to call a function
+    if tool_calls:
+        # Step 3: call the function
+        # Note: the JSON response may not always be valid; be sure to handle errors
+        available_functions = {
+            "run_cypher_query": run_cypher_query,
+        }  # only one function in this example, but you can have multiple
+        messages.append(response_message)  # extend conversation with assistant's reply
+        # Step 4: send the info for each function call and function response to the model
+        for tool_call in tool_calls:
+            function_name = tool_call.function.name
+            function_to_call = available_functions[function_name]
+            function_args = json.loads(tool_call.function.arguments)
+            q = function_args.get("query")
+            print(f"OpenAI suggested query: {q}")
+            function_response = function_to_call(g, q)
+            messages.append(
+                {
+                    "tool_call_id": tool_call.id,
+                    "role": "tool",
+                    "name": function_name,
+                    "content": str(function_response),
+                }
+            )  # extend conversation with function response
+        second_response = client.chat.completions.create(
+            model=model,
+            messages=messages,
+        )  # get a new response from the model where it can see the function response
+
+        return second_response.choices[0].message.content
+
+    return None
+
+def query(kg, question:str, messages:list, model="gpt-4-1106-preview") -> str:
+    # Build system message
+    if len(messages) == 0:
+        # graph schema
+        schema_desc = _graph_schema_to_prompt(kg.schema)
+        messages.append(
+            {"role": "system", "content": f"""You are a graph database expert fluent in the Cypher query language
+                with access to the following knowledge graph: {schema_desc}"""})
+
+    # Get an answer
+    answer = run_conversation(kg.graph, question, messages, model)
+
+    # Append answer to conversation
+    messages.append({"role": "assistant", "content": answer})
+
+    return answer

graphrag_sdk-0.1.0/graphrag_sdk/schema/__init__.py

@@ -0,0 +1 @@
+from .schema import Schema