use KnowledgeGraph typed dict for graph API response

lightrag/types.py

@@ -1,7 +1,26 @@
 from pydantic import BaseModel
-from typing import List
+from typing import List, Dict, Any
 
 
 class GPTKeywordExtractionFormat(BaseModel):
     high_level_keywords: List[str]
     low_level_keywords: List[str]
+
+
+class KnowledgeGraphNode(BaseModel):
+    id: str
+    labels: List[str]
+    properties: Dict[str, Any]  # anything else goes here
+
+
+class KnowledgeGraphEdge(BaseModel):
+    id: str
+    type: str
+    source: str  # id of source node
+    target: str  # id of target node
+    properties: Dict[str, Any]  # anything else goes here
+
+
+class KnowledgeGraph(BaseModel):
+    nodes: List[KnowledgeGraphNode] = []
+    edges: List[KnowledgeGraphEdge] = []