'commit'

AI/Neo4j/KnowledgeGraph.py

@@ -7,8 +7,8 @@ from Config import API_KEY, MODEL_R1, MODEL_URL
 
 
 class KnowledgeGraph:
-    def __init__(self, shiti_content: str):
+    def __init__(self, content: str):
-        self.shiti_content = shiti_content
+        self.content = content
         self.client = OpenAI(api_key=API_KEY, base_url=MODEL_URL)
 
     def _generate_stream(self) -> Iterator[ChatCompletionChunk]:
@@ -19,7 +19,7 @@ class KnowledgeGraph:
             model=MODEL_R1,
             messages=[
                 {"role": "system", "content": system_prompt},
-                {"role": "user", "content": self.shiti_content}
+                {"role": "user", "content": self.content}
             ],
             stream=True,
             timeout=300