mirror of
https://github.com/hwchase17/langchain
synced 2024-11-08 07:10:35 +00:00
ebf998acb6
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com> Co-authored-by: Lance Martin <lance@langchain.dev> Co-authored-by: Jacob Lee <jacoblee93@gmail.com>
74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
from langchain.graphs.graph_document import (
|
|
Node as BaseNode,
|
|
Relationship as BaseRelationship,
|
|
)
|
|
from typing import List, Optional
|
|
from langchain.pydantic_v1 import Field, BaseModel
|
|
|
|
|
|
class Property(BaseModel):
|
|
"""A single property consisting of key and value"""
|
|
|
|
key: str = Field(..., description="key")
|
|
value: str = Field(..., description="value")
|
|
|
|
|
|
class Node(BaseNode):
|
|
properties: Optional[List[Property]] = Field(
|
|
None, description="List of node properties"
|
|
)
|
|
|
|
|
|
class Relationship(BaseRelationship):
|
|
properties: Optional[List[Property]] = Field(
|
|
None, description="List of relationship properties"
|
|
)
|
|
|
|
|
|
class KnowledgeGraph(BaseModel):
|
|
"""Generate a knowledge graph with entities and relationships."""
|
|
|
|
nodes: List[Node] = Field(..., description="List of nodes in the knowledge graph")
|
|
rels: List[Relationship] = Field(
|
|
..., description="List of relationships in the knowledge graph"
|
|
)
|
|
|
|
|
|
def format_property_key(s: str) -> str:
|
|
words = s.split()
|
|
if not words:
|
|
return s
|
|
first_word = words[0].lower()
|
|
capitalized_words = [word.capitalize() for word in words[1:]]
|
|
return "".join([first_word] + capitalized_words)
|
|
|
|
|
|
def props_to_dict(props) -> dict:
|
|
"""Convert properties to a dictionary."""
|
|
properties = {}
|
|
if not props:
|
|
return properties
|
|
for p in props:
|
|
properties[format_property_key(p.key)] = p.value
|
|
return properties
|
|
|
|
|
|
def map_to_base_node(node: Node) -> BaseNode:
|
|
"""Map the KnowledgeGraph Node to the base Node."""
|
|
properties = props_to_dict(node.properties) if node.properties else {}
|
|
# Add name property for better Cypher statement generation
|
|
properties["name"] = node.id.title()
|
|
return BaseNode(
|
|
id=node.id.title(), type=node.type.capitalize(), properties=properties
|
|
)
|
|
|
|
|
|
def map_to_base_relationship(rel: Relationship) -> BaseRelationship:
|
|
"""Map the KnowledgeGraph Relationship to the base Relationship."""
|
|
source = map_to_base_node(rel.source)
|
|
target = map_to_base_node(rel.target)
|
|
properties = props_to_dict(rel.properties) if rel.properties else {}
|
|
return BaseRelationship(
|
|
source=source, target=target, type=rel.type, properties=properties
|
|
)
|