diff --git a/AI/Neo4j/KnowledgeGraph.py b/AI/Neo4j/KnowledgeGraph.py index 61e6fb18..5bb4034a 100644 --- a/AI/Neo4j/KnowledgeGraph.py +++ b/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): - self.shiti_content = shiti_content + def __init__(self, content: str): + 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