feat: solution for 'Повторный экзамен: Граф с рефлексией и доработкой'

src/nodes.py

@@ -1,21 +1,83 @@
-from langchain_core.messages import HumanMessage, AIMessage
-from typing import Dict, Any
+"""
+Node definitions for the graph.
+Includes base Node, ReflectionNode, RewritingNode, InputNode, and OutputNode.
+"""
 
-def generate_response(state: Dict[str, Any]) -> Dict[str, Any]:
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from .llm_integration import get_llm
+
+
+class BaseNode(ABC):
     """
-    Simple node that echoes the user's message as an AI response.
+    Abstract base class for all nodes in the graph.
+    Each node must implement the `process` method.
     """
-    messages = state.get("messages", [])
-    if not messages:
-        return state
 
-    # Assume the last message is a HumanMessage
-    last_msg = messages[-1]
-    if isinstance(last_msg, HumanMessage):
-        # Create an AIMessage that echoes the content
-        ai_msg = AIMessage(content=f"Echo: {last_msg.content}")
-        messages.append(ai_msg)
+    def __init__(self, node_id: str):
+        self.node_id = node_id
 
-    # Update the state with the new messages list
-    state["messages"] = messages
-    return state
+    @abstractmethod
+    def process(self, input_data: Any) -> Any:
+        """
+        Process the input data and return the output.
+        """
+        pass
+
+
+class InputNode(BaseNode):
+    """
+    Node that simply passes through the input data.
+    """
+
+    def process(self, input_data: Any) -> Any:
+        return input_data
+
+
+class OutputNode(BaseNode):
+    """
+    Node that collects the final output.
+    """
+
+    def process(self, input_data: Any) -> Any:
+        return input_data
+
+
+class ReflectionNode(BaseNode):
+    """
+    Node that generates reflective insights from the input text using an LLM.
+    """
+
+    def __init__(self, node_id: str, prompt_template: str = None):
+        super().__init__(node_id)
+        self.prompt_template = (
+            prompt_template
+            or "Please reflect on the following text:\n\n{input_text}\n\nReflection:"
+        )
+        self.llm = get_llm()
+
+    def process(self, input_data: str) -> Dict[str, str]:
+        prompt = self.prompt_template.format(input_text=input_data)
+        reflection = self.llm(prompt)
+        return {"reflection": reflection.strip()}
+
+
+class RewritingNode(BaseNode):
+    """
+    Node that rewrites the input text according to a specified style or instruction.
+    """
+
+    def __init__(self, node_id: str, style: str = "formal"):
+        super().__init__(node_id)
+        self.style = style
+        self.llm = get_llm()
+
+    def process(self, input_data: Dict[str, str]) -> Dict[str, str]:
+        # Expecting input_data to contain 'reflection' key
+        reflection = input_data.get("reflection", "")
+        prompt = (
+            f"Rewrite the following reflection in a {self.style} style:\n\n{reflection}\n\nRewritten:"
+        )
+        rewritten = self.llm(prompt)
+        return {"rewritten": rewritten.strip()}