use logging

kg_builder/src/knowledge_graph_builder.py

@@ -1,3 +1,4 @@
+import logging
 
 from api_connections import get_graph_connection
 
@@ -13,15 +14,25 @@ from models import Node, Relationship, KnowledgeGraph
 from utils import map_to_base_node, map_to_base_relationship
 from api_connections import get_extraction_chain
 
+logging.basicConfig(format='%(name)s - %(levelname)s - %(message)s', level=logging.INFO)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
 def extract_and_store_graph(
     document: Document,
-    category: str,
+    data_source_name: str,
     nodes:Optional[List[str]] = None,
     rels:Optional[List[str]]=None) -> None:
-    
-    graph = get_graph_connection(category)
-    # Extract graph data using OpenAI functions
-    extract_chain = get_extraction_chain(category, nodes, rels)
+    """
+    Extract data from text document and add nodes and relationships to knowledge graph
+    :param document: Text document
+    :param data_source_Name: Data source name, e.g. "Traffic Law"
+    :param nodes: TODO
+    :param rels: TODO
+    """
+
+    logger.info("Extract graph data using OpenAI functions ...")
+    extract_chain = get_extraction_chain(data_source_name, nodes, rels)
     data = extract_chain.invoke(document.page_content)['function']
     # Construct a graph document
     graph_document = GraphDocument(
@@ -29,6 +40,8 @@ def extract_and_store_graph(
       relationships = [map_to_base_relationship(rel) for rel in data.rels],
       source = document
     )
+
     # Store information into a graph
+    graph = get_graph_connection(data_source_name)
     graph.add_graph_documents([graph_document])