feat: solution for 'Повторный экзамен: Граф с рефлексией и доработкой'
This commit is contained in:
+62
-28
@@ -1,39 +1,73 @@
|
||||
"""
|
||||
Graph definition for the LangGraph workflow.
|
||||
|
||||
The graph consists of two nodes:
|
||||
1. ReflectNode – generates a reflection of the user input.
|
||||
2. RewriteNode – rewrites the reflection into a more formal style.
|
||||
|
||||
The graph starts at the reflect node, then proceeds to the rewrite node,
|
||||
and finishes with the rewritten output.
|
||||
Graph implementation that connects nodes and executes them in sequence.
|
||||
"""
|
||||
|
||||
from langgraph.graph import StateGraph
|
||||
from src.nodes.reflect import ReflectNode
|
||||
from src.nodes.rewrite import RewriteNode
|
||||
from typing import Dict, List
|
||||
|
||||
from .nodes import BaseNode, InputNode, OutputNode, ReflectionNode, RewritingNode
|
||||
|
||||
|
||||
def build_graph():
|
||||
class Graph:
|
||||
"""
|
||||
Build and compile the LangGraph graph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
langgraph.graph.Graph
|
||||
The compiled graph ready for invocation.
|
||||
Simple directed acyclic graph for node execution.
|
||||
"""
|
||||
builder = StateGraph()
|
||||
builder.add_node("reflect", ReflectNode.run)
|
||||
builder.add_node("rewrite", RewriteNode.run)
|
||||
|
||||
# Entry point is the reflect node
|
||||
builder.set_entry_point("reflect")
|
||||
def __init__(self):
|
||||
self.nodes: Dict[str, BaseNode] = {}
|
||||
self.edges: Dict[str, List[str]] = {}
|
||||
|
||||
# Define the flow: reflect -> rewrite
|
||||
builder.add_edge("reflect", "rewrite")
|
||||
def add_node(self, node: BaseNode):
|
||||
self.nodes[node.node_id] = node
|
||||
self.edges.setdefault(node.node_id, [])
|
||||
|
||||
# Finish at the rewrite node
|
||||
builder.set_finish("rewrite")
|
||||
def add_edge(self, from_node_id: str, to_node_id: str):
|
||||
if from_node_id not in self.nodes or to_node_id not in self.nodes:
|
||||
raise ValueError("Both nodes must be added before creating an edge.")
|
||||
self.edges[from_node_id].append(to_node_id)
|
||||
|
||||
return builder.compile()
|
||||
def _find_start_node(self) -> str:
|
||||
# Node with no incoming edges
|
||||
all_targets = {t for targets in self.edges.values() for t in targets}
|
||||
for node_id in self.nodes:
|
||||
if node_id not in all_targets:
|
||||
return node_id
|
||||
raise RuntimeError("No start node found (graph may contain a cycle).")
|
||||
|
||||
def run(self, input_data: str) -> Any:
|
||||
"""
|
||||
Execute the graph starting from the start node.
|
||||
"""
|
||||
current_node_id = self._find_start_node()
|
||||
data = input_data
|
||||
while True:
|
||||
node = self.nodes[current_node_id]
|
||||
data = node.process(data)
|
||||
successors = self.edges.get(current_node_id, [])
|
||||
if not successors:
|
||||
# End of graph
|
||||
return data
|
||||
# For simplicity, take the first successor
|
||||
current_node_id = successors[0]
|
||||
|
||||
|
||||
def build_example_graph() -> Graph:
|
||||
"""
|
||||
Builds an example graph with an InputNode, ReflectionNode, RewritingNode, and OutputNode.
|
||||
"""
|
||||
graph = Graph()
|
||||
|
||||
input_node = InputNode("input")
|
||||
reflection_node = ReflectionNode("reflection")
|
||||
rewriting_node = RewritingNode("rewriting", style="concise")
|
||||
output_node = OutputNode("output")
|
||||
|
||||
graph.add_node(input_node)
|
||||
graph.add_node(reflection_node)
|
||||
graph.add_node(rewriting_node)
|
||||
graph.add_node(output_node)
|
||||
|
||||
graph.add_edge("input", "reflection")
|
||||
graph.add_edge("reflection", "rewriting")
|
||||
graph.add_edge("rewriting", "output")
|
||||
|
||||
return graph
|
||||
@@ -0,0 +1,33 @@
|
||||
"""
|
||||
LLM integration module for LangChain with support for OpenAI and Ollama.
|
||||
Provides a reusable LLM client based on environment configuration.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
from langchain.llms import OpenAI, Ollama
|
||||
from langchain.chat_models import ChatOpenAI, ChatOllama
|
||||
|
||||
# Environment variable to select provider: "openai" or "ollama"
|
||||
LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").lower()
|
||||
|
||||
|
||||
def get_llm() -> Union[OpenAI, Ollama, ChatOpenAI, ChatOllama]:
|
||||
"""
|
||||
Returns an LLM instance based on the configured provider.
|
||||
|
||||
For OpenAI, uses the default OpenAI LLM (text-davinci-003 or gpt-3.5-turbo).
|
||||
For Ollama, uses the default Ollama LLM (e.g., llama2).
|
||||
|
||||
Raises:
|
||||
ValueError: If an unsupported provider is specified.
|
||||
"""
|
||||
if LLM_PROVIDER == "openai":
|
||||
# Use ChatOpenAI for GPT-3.5-turbo by default
|
||||
return ChatOpenAI(temperature=0.7)
|
||||
elif LLM_PROVIDER == "ollama":
|
||||
# Use ChatOllama for local models
|
||||
return ChatOllama(model="llama2", temperature=0.7)
|
||||
else:
|
||||
raise ValueError(f"Unsupported LLM provider: {LLM_PROVIDER}")
|
||||
+28
-9
@@ -1,18 +1,37 @@
|
||||
"""
|
||||
Entry point for running the LangGraph workflow.
|
||||
|
||||
This script demonstrates how to invoke the graph with a sample input.
|
||||
Entry point for running the graph with user-provided text.
|
||||
"""
|
||||
|
||||
from src.graph import build_graph
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
from .graph import build_example_graph
|
||||
|
||||
|
||||
def main():
|
||||
graph = build_graph()
|
||||
# Sample input
|
||||
input_state = {"input": "Hello world"}
|
||||
result = graph.invoke(input_state)
|
||||
print("Graph output:", result)
|
||||
parser = argparse.ArgumentParser(description="Run the reflection and rewriting graph.")
|
||||
parser.add_argument(
|
||||
"text",
|
||||
nargs="?",
|
||||
help="Input text to process. If omitted, reads from stdin.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.text:
|
||||
input_text = args.text
|
||||
else:
|
||||
input_text = sys.stdin.read()
|
||||
|
||||
graph = build_example_graph()
|
||||
result = graph.run(input_text)
|
||||
|
||||
# The final node returns a dict with 'rewritten' key
|
||||
if isinstance(result, dict) and "rewritten" in result:
|
||||
print("Rewritten Text:\n")
|
||||
print(result["rewritten"])
|
||||
else:
|
||||
print("Result:")
|
||||
print(result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+78
-16
@@ -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()}
|
||||
Reference in New Issue
Block a user