83 lines
2.2 KiB
Python
83 lines
2.2 KiB
Python
"""
|
|
Node definitions for the graph.
|
|
Includes base Node, ReflectionNode, RewritingNode, InputNode, and OutputNode.
|
|
"""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any, Dict
|
|
|
|
from .llm_integration import get_llm
|
|
|
|
|
|
class BaseNode(ABC):
|
|
"""
|
|
Abstract base class for all nodes in the graph.
|
|
Each node must implement the `process` method.
|
|
"""
|
|
|
|
def __init__(self, node_id: str):
|
|
self.node_id = node_id
|
|
|
|
@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()} |