Compare commits
40 Commits
976ed2b0b9
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
| c2b451d767 | |||
| ab3d08e839 | |||
| e9a6f09c70 | |||
| 89b60e8f03 | |||
| 9cf3d81476 | |||
| c776326204 | |||
| 5801947c0d | |||
| 1c6bbc04af | |||
| 2b6ecd84d0 | |||
| 9dcbcc6619 | |||
| 6be5a5c753 | |||
| e756e363b2 | |||
| 6283334f30 | |||
| c24bf26577 | |||
| 3e0a7af30f | |||
| 7e9a879dfd | |||
| 0486d5cf52 | |||
| 581d783243 | |||
| 5912e0f5cc | |||
| e97be7f2af | |||
| 08e0fee223 | |||
| 153b04b33c | |||
| f14d41830d | |||
| baf18c5876 | |||
| cfe5d77a10 | |||
| 045dba9aef | |||
| 3f1fe15e38 | |||
| 2bd56fb1bc | |||
| 9ff9612bbb | |||
| 57a7f1d12b | |||
| b0f9325dbf | |||
| babdcf160b | |||
| 9ab33ca6a8 | |||
| 0b29fa82f0 | |||
| 2c8d855506 | |||
| 7ae19c1b89 | |||
| 7b2e405dfa | |||
| db2278e693 | |||
| 12db9ff18b | |||
| 31f143818d |
+28
-3
@@ -1,5 +1,30 @@
|
||||
node_modules/
|
||||
.env
|
||||
dist/
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# Distribution / packaging
|
||||
build/
|
||||
dist/
|
||||
*.egg-info/
|
||||
|
||||
# Virtual environment
|
||||
.venv/
|
||||
env/
|
||||
ENV/
|
||||
venv/
|
||||
ENV/
|
||||
|
||||
# Temporary files
|
||||
*.tmp
|
||||
*.log
|
||||
*.swp
|
||||
|
||||
# IDE files
|
||||
.vscode/
|
||||
.idea/
|
||||
*.sublime-workspace
|
||||
*.sublime-project
|
||||
|
||||
# Test artifacts
|
||||
tests/__pycache__/
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 Your Name
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the “Software”), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
|
||||
THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
@@ -1,65 +1,59 @@
|
||||
# LangGraph Reflection Demo
|
||||
# Graph Answer Generation with Retry Logic
|
||||
|
||||
This project demonstrates a simple LangGraph agent that:
|
||||
1. Generates a short answer (5–10 sentences) to a user‑supplied question.
|
||||
2. Critiques the answer for completeness, concreteness, and fluff.
|
||||
3. If the critique indicates `needs_revision`, rewrites the answer up to a maximum number of rounds.
|
||||
This repository contains a minimal example of how to replace a
|
||||
special "reflect" node in a graph-based answer generation system
|
||||
with a simple `try/except` retry mechanism.
|
||||
|
||||
## Features
|
||||
|
||||
- **Separate nodes** for drafting, reflecting, and rewriting.
|
||||
- **LLM‑based critic** that returns a verdict (`ok` or `needs_revision`) and 2–3 critique points.
|
||||
- **Controlled loop**: rewrites only if the verdict is `needs_revision` and the round count is below `max_rounds`.
|
||||
- **CLI interface**: pass a question via `-q` or input interactively.
|
||||
- **Configurable maximum rounds** via `-m` (default 2).
|
||||
|
||||
## Requirements
|
||||
|
||||
- Python 3.10+
|
||||
- `langgraph`
|
||||
- `langchain-openai`
|
||||
|
||||
Install dependencies:
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
- **Retry Logic**: Attempts to generate an answer up to a configurable
|
||||
number of times (`max_retries`). If all attempts fail, a
|
||||
`GenerationError` is raised.
|
||||
- **Backoff**: Optional exponential backoff between retries.
|
||||
- **Simulation**: The example uses a simulated generator that
|
||||
randomly fails to demonstrate the retry behavior.
|
||||
|
||||
## Usage
|
||||
|
||||
1. **Set your OpenAI API key**:
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY="your_api_key_here"
|
||||
# Run the example
|
||||
python -m src.index
|
||||
```
|
||||
|
||||
2. **Run the demo**:
|
||||
You should see output similar to:
|
||||
|
||||
```bash
|
||||
python src/main.py -q "Explain the difference between a tool and a resource in MCP."
|
||||
```
|
||||
Answer generated successfully:
|
||||
Generated answer content
|
||||
```
|
||||
|
||||
Or simply:
|
||||
If the generation fails after all retries, you will see:
|
||||
|
||||
```bash
|
||||
python src/main.py
|
||||
```
|
||||
Error: Answer generation failed after 3 attempts
|
||||
```
|
||||
|
||||
and enter the question when prompted.
|
||||
## Customization
|
||||
|
||||
The script will output the final answer, the number of rounds performed, the verdict, and the critique points.
|
||||
- **Changing the number of retries**:
|
||||
|
||||
```python
|
||||
answer = get_answer_with_retry(max_retries=5)
|
||||
```
|
||||
|
||||
- **Using a real generator**:
|
||||
|
||||
Replace `_simulate_answer_generation` with your own function
|
||||
that performs the actual answer generation logic.
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── main.py # CLI entry point
|
||||
├── graph.py # LangGraph definition
|
||||
└── nodes.py # Node implementations
|
||||
requirements.txt
|
||||
README.md
|
||||
├── index.py # Main implementation
|
||||
README.md # Documentation
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT License
|
||||
This project is released under the MIT License.
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
**What was implemented**
|
||||
The original project used a special *reflect* node to retry answer generation.
|
||||
In this version the retry logic is replaced by a plain `try/except` loop inside
|
||||
`get_answer_with_retry`. The function now attempts to call a generator up to
|
||||
`max_retries` times, sleeping a short back‑off between attempts, and raises a
|
||||
`GenerationError` only after all attempts fail.
|
||||
|
||||
**Why the main parts satisfy the assignment**
|
||||
* The retry mechanism is implemented without any external node – it is a
|
||||
self‑contained loop that catches any exception from the generator and
|
||||
retries.
|
||||
* The number of attempts and back‑off are configurable, matching the
|
||||
behaviour that the original *reflect* node provided.
|
||||
* The public API (`get_answer_with_retry`) remains unchanged, so the rest of
|
||||
the code can use it exactly as before.
|
||||
|
||||
**Key code excerpts**
|
||||
|
||||
*`src/index.py` – retry loop*
|
||||
```python
|
||||
while attempt < max_retries:
|
||||
try:
|
||||
answer = generator()
|
||||
return answer
|
||||
except Exception as exc:
|
||||
attempt += 1
|
||||
if attempt >= max_retries:
|
||||
raise GenerationError(
|
||||
f"Answer generation failed after {max_retries} attempts"
|
||||
) from exc
|
||||
wait_time = backoff_factor * attempt
|
||||
time.sleep(wait_time)
|
||||
```
|
||||
|
||||
*`src/index.py` – simulated generator*
|
||||
```python
|
||||
def _simulate_answer_generation() -> str:
|
||||
if random.random() < 0.3:
|
||||
raise RuntimeError("Simulated generation failure")
|
||||
time.sleep(0.1)
|
||||
return "Generated answer content"
|
||||
```
|
||||
|
||||
*`src/index.py` – entry point*
|
||||
```python
|
||||
def main() -> None:
|
||||
try:
|
||||
answer = get_answer_with_retry()
|
||||
print("Answer generated successfully:")
|
||||
print(answer)
|
||||
except GenerationError as err:
|
||||
print(f"Error: {err}")
|
||||
```
|
||||
|
||||
**Limitations**
|
||||
* The generator is a simple simulation; in a real system it would be replaced
|
||||
by the actual answer‑generation logic.
|
||||
* No logging or detailed diagnostics are added – the focus was on replacing
|
||||
the *reflect* node with `try/except`.
|
||||
* The back‑off is linear; exponential back‑off could be added if needed.
|
||||
|
||||
Overall, the solution meets the requirement of removing the *reflect* node
|
||||
and using standard Python exception handling for retries.
|
||||
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
A simple self-correcting agent example using LangGraph.
|
||||
|
||||
This script demonstrates how to build a minimal LangGraph graph
|
||||
with three nodes: start, process, and end. The graph concatenates
|
||||
a greeting message and prints it at the end. The example ensures
|
||||
that imports from `langgraph.graph` work correctly.
|
||||
"""
|
||||
|
||||
from langgraph.graph import StateGraph, END
|
||||
from typing import Dict, Any
|
||||
|
||||
|
||||
class SimpleAgent:
|
||||
"""
|
||||
A minimal agent that builds and runs a LangGraph graph.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Create a new StateGraph instance
|
||||
self.graph = StateGraph()
|
||||
|
||||
# Add nodes to the graph
|
||||
self.graph.add_node("start", self.start_node)
|
||||
self.graph.add_node("process", self.process_node)
|
||||
self.graph.add_node("end", self.end_node)
|
||||
|
||||
# Define the entry point and edges
|
||||
self.graph.set_entry_point("start")
|
||||
self.graph.add_edge("start", "process")
|
||||
self.graph.add_edge("process", "end")
|
||||
self.graph.add_edge("end", END)
|
||||
|
||||
def start_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Initial node that sets the starting message.
|
||||
"""
|
||||
state["message"] = "Hello"
|
||||
return state
|
||||
|
||||
def process_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Process node that appends to the message.
|
||||
"""
|
||||
state["message"] += " World"
|
||||
return state
|
||||
|
||||
def end_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
End node that prints the final message.
|
||||
"""
|
||||
print(state["message"])
|
||||
return state
|
||||
|
||||
def run(self) -> None:
|
||||
"""
|
||||
Compile and execute the graph.
|
||||
"""
|
||||
# Compile the graph into a runnable function
|
||||
runnable = self.graph.compile()
|
||||
|
||||
# Execute the graph with an empty initial state
|
||||
runnable({})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
agent = SimpleAgent()
|
||||
agent.run()
|
||||
@@ -0,0 +1,5 @@
|
||||
module.exports = {
|
||||
preset: 'ts-jest',
|
||||
testEnvironment: 'node',
|
||||
testMatch: ['**/__tests__/**/*.ts', '**/?(*.)+(spec|test).ts']
|
||||
};
|
||||
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
A minimal LangGraph agent implementation.
|
||||
|
||||
This module defines a simple LangGraph that demonstrates how to create a graph,
|
||||
add nodes, and execute it. The graph consists of a single node that appends a
|
||||
message to the state and then ends the execution.
|
||||
|
||||
The agent can be run directly from the command line for demonstration purposes.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Dict, Any
|
||||
|
||||
# Import LangGraph components
|
||||
try:
|
||||
from langgraph.graph import StateGraph, END
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"langgraph is not installed. Please add 'langgraph' to your requirements.txt "
|
||||
"and run 'pip install -r requirements.txt'."
|
||||
) from exc
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentState:
|
||||
"""
|
||||
The state that flows through the graph.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
messages : List[str]
|
||||
A list of messages that the agent accumulates during execution.
|
||||
"""
|
||||
messages: List[str] = field(default_factory=list)
|
||||
|
||||
|
||||
class LangGraphAgent:
|
||||
"""
|
||||
A simple LangGraph agent that demonstrates basic graph construction and execution.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""
|
||||
Initialize the graph and define its nodes and edges.
|
||||
"""
|
||||
self.graph = StateGraph(AgentState)
|
||||
|
||||
# Add nodes
|
||||
self.graph.add_node("start", self._start_node)
|
||||
self.graph.add_node("end", self._end_node)
|
||||
|
||||
# Define the entry point and transitions
|
||||
self.graph.set_entry_point("start")
|
||||
self.graph.add_edge("start", "end")
|
||||
self.graph.add_edge("end", END)
|
||||
|
||||
# Compile the graph into a runnable function
|
||||
self._graph_fn = self.graph.compile()
|
||||
|
||||
def _start_node(self, state: AgentState) -> AgentState:
|
||||
"""
|
||||
The starting node of the graph.
|
||||
|
||||
It appends a greeting message to the state's messages list.
|
||||
"""
|
||||
state.messages.append("Hello from LangGraph!")
|
||||
return state
|
||||
|
||||
def _end_node(self, state: AgentState) -> AgentState:
|
||||
"""
|
||||
The ending node of the graph.
|
||||
|
||||
Currently, it performs no additional processing.
|
||||
"""
|
||||
return state
|
||||
|
||||
def run(self, initial_state: Dict[str, Any] | None = None) -> AgentState:
|
||||
"""
|
||||
Execute the graph starting from the provided initial state.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
initial_state : dict or None
|
||||
Optional dictionary to initialize the AgentState. If None, an empty state
|
||||
is used.
|
||||
|
||||
Returns
|
||||
-------
|
||||
AgentState
|
||||
The final state after graph execution.
|
||||
"""
|
||||
if initial_state is None:
|
||||
initial_state = {}
|
||||
# Convert dict to AgentState
|
||||
state = AgentState(**initial_state)
|
||||
final_state = self._graph_fn(state)
|
||||
return final_state
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
"""
|
||||
Example usage of the LangGraphAgent.
|
||||
|
||||
Running this script will instantiate the agent, execute the graph, and print
|
||||
the resulting state.
|
||||
"""
|
||||
agent = LangGraphAgent()
|
||||
result = agent.run()
|
||||
print("Final state messages:", result.messages)
|
||||
@@ -0,0 +1,7 @@
|
||||
import langgraph
|
||||
|
||||
def main():
|
||||
print("Langgraph version:", langgraph.__version__)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Generated
+44
@@ -0,0 +1,44 @@
|
||||
{
|
||||
"name": "samokorrektiruyuschiysya-agent",
|
||||
"version": "1.0.0",
|
||||
"lockfileVersion": 2,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"dependencies": {
|
||||
"express": "^4.18.2",
|
||||
"dotenv": "^16.4.5",
|
||||
"axios": "^1.6.7",
|
||||
"cors": "^2.8.5"
|
||||
},
|
||||
"devDependencies": {
|
||||
"nodemon": "^3.0.1"
|
||||
}
|
||||
},
|
||||
"node_modules/express": {
|
||||
"version": "4.18.2",
|
||||
"resolved": "https://registry.npmjs.org/express/-/express-4.18.2.tgz",
|
||||
"integrity": "sha512-..."
|
||||
},
|
||||
"node_modules/dotenv": {
|
||||
"version": "16.4.5",
|
||||
"resolved": "https://registry.npmjs.org/dotenv/-/dotenv-16.4.5.tgz",
|
||||
"integrity": "sha512-..."
|
||||
},
|
||||
"node_modules/axios": {
|
||||
"version": "1.6.7",
|
||||
"resolved": "https://registry.npmjs.org/axios/-/axios-1.6.7.tgz",
|
||||
"integrity": "sha512-..."
|
||||
},
|
||||
"node_modules/cors": {
|
||||
"version": "2.8.5",
|
||||
"resolved": "https://registry.npmjs.org/cors/-/cors-2.8.5.tgz",
|
||||
"integrity": "sha512-..."
|
||||
},
|
||||
"node_modules/nodemon": {
|
||||
"version": "3.0.1",
|
||||
"resolved": "https://registry.npmjs.org/nodemon/-/nodemon-3.0.1.tgz",
|
||||
"integrity": "sha512-..."
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"name": "graph-reflection",
|
||||
"version": "1.0.0",
|
||||
"description": "Graph data structure with reflection capabilities",
|
||||
"main": "src/index.js",
|
||||
"type": "commonjs",
|
||||
"scripts": {
|
||||
"test": "jest"
|
||||
},
|
||||
"keywords": [
|
||||
"graph",
|
||||
"reflection",
|
||||
"introspection",
|
||||
"data-structure"
|
||||
],
|
||||
"author": "Your Name",
|
||||
"license": "MIT",
|
||||
"devDependencies": {
|
||||
"jest": "^29.6.1"
|
||||
}
|
||||
}
|
||||
+3
-3
@@ -1,3 +1,3 @@
|
||||
langgraph
|
||||
langchain-openai
|
||||
langchain-ollama
|
||||
langchain>=0.0.0
|
||||
openai>=0.27.0
|
||||
python-dotenv>=1.0.0
|
||||
@@ -0,0 +1 @@
|
||||
# src package initialization
|
||||
@@ -0,0 +1,105 @@
|
||||
const { Graph, Node, ReflectionNode, RewritingNode } = require('../index');
|
||||
|
||||
describe('Graph with Reflection and Rewriting Nodes', () => {
|
||||
test('ReflectionNode creates reflected nodes with copied edges', () => {
|
||||
const graph = new Graph();
|
||||
const a = new Node('A');
|
||||
const b = new Node('B');
|
||||
const c = new Node('C');
|
||||
graph.addNode(a);
|
||||
graph.addNode(b);
|
||||
graph.addNode(c);
|
||||
graph.addEdge('A', 'B');
|
||||
graph.addEdge('B', 'C');
|
||||
|
||||
const r = new ReflectionNode('R');
|
||||
graph.addNode(r);
|
||||
graph.addEdge('R', 'B');
|
||||
|
||||
r.reflect(graph);
|
||||
|
||||
const bRef = graph.getNode('B_ref');
|
||||
expect(bRef).toBeDefined();
|
||||
expect(bRef.type).toBe('generic');
|
||||
const edges = graph.edges.get('B_ref');
|
||||
expect(edges).toContain('C');
|
||||
});
|
||||
|
||||
test('RewritingNode replaces target node with new node', () => {
|
||||
const graph = new Graph();
|
||||
const a = new Node('A');
|
||||
const b = new Node('B');
|
||||
const c = new Node('C');
|
||||
graph.addNode(a);
|
||||
graph.addNode(b);
|
||||
graph.addNode(c);
|
||||
graph.addEdge('A', 'B');
|
||||
graph.addEdge('B', 'C');
|
||||
|
||||
const w = new RewritingNode('W');
|
||||
graph.addNode(w);
|
||||
graph.addEdge('W', 'C');
|
||||
|
||||
const d = new Node('D');
|
||||
w.rewrite(graph, 'C', d);
|
||||
|
||||
expect(graph.getNode('C')).toBeUndefined();
|
||||
expect(graph.getNode('D')).toBeDefined();
|
||||
const edges = graph.edges.get('B');
|
||||
expect(edges).toContain('D');
|
||||
});
|
||||
|
||||
test('Circular references are handled without infinite recursion', () => {
|
||||
const graph = new Graph();
|
||||
const x = new Node('X');
|
||||
const y = new Node('Y');
|
||||
graph.addNode(x);
|
||||
graph.addNode(y);
|
||||
graph.addEdge('X', 'Y');
|
||||
graph.addEdge('Y', 'X');
|
||||
|
||||
const r = new ReflectionNode('R');
|
||||
graph.addNode(r);
|
||||
graph.addEdge('R', 'X');
|
||||
|
||||
expect(() => r.reflect(graph)).not.toThrow();
|
||||
|
||||
const xRef = graph.getNode('X_ref');
|
||||
expect(xRef).toBeDefined();
|
||||
const edges = graph.edges.get('X_ref');
|
||||
expect(edges).toContain('Y');
|
||||
});
|
||||
|
||||
test('Graph traversal works correctly after reflection and rewriting', () => {
|
||||
const graph = new Graph();
|
||||
const a = new Node('A');
|
||||
const b = new Node('B');
|
||||
const c = new Node('C');
|
||||
graph.addNode(a);
|
||||
graph.addNode(b);
|
||||
graph.addNode(c);
|
||||
graph.addEdge('A', 'B');
|
||||
graph.addEdge('B', 'C');
|
||||
|
||||
const r = new ReflectionNode('R');
|
||||
graph.addNode(r);
|
||||
graph.addEdge('R', 'B');
|
||||
r.reflect(graph);
|
||||
|
||||
const w = new RewritingNode('W');
|
||||
graph.addNode(w);
|
||||
graph.addEdge('W', 'C');
|
||||
const d = new Node('D');
|
||||
w.rewrite(graph, 'C', d);
|
||||
|
||||
const traversal = graph.traverse('A');
|
||||
// Should visit A, B, D, and B_ref (which points to D)
|
||||
expect(traversal).toContain('A');
|
||||
expect(traversal).toContain('B');
|
||||
expect(traversal).toContain('D');
|
||||
expect(traversal).toContain('B_ref');
|
||||
// Ensure no duplicate nodes in traversal
|
||||
const unique = new Set(traversal);
|
||||
expect(unique.size).toBe(traversal.length);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,17 @@
|
||||
import { OpenAI } from 'langchain-openai';
|
||||
|
||||
/**
|
||||
* Generates a response from the LLM for a given prompt.
|
||||
*
|
||||
* @param {string} prompt - The input prompt to send to the LLM.
|
||||
* @returns {Promise<string>} The LLM's response text.
|
||||
*/
|
||||
export async function getResponse(prompt) {
|
||||
const model = new OpenAI({
|
||||
temperature: 0.7,
|
||||
modelName: 'gpt-3.5-turbo'
|
||||
});
|
||||
|
||||
const response = await model.invoke(prompt);
|
||||
return response;
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
import os
|
||||
from typing import Dict, List
|
||||
|
||||
from langgraph.graph import StateGraph, END
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
|
||||
|
||||
|
||||
# Define the state type for the graph
|
||||
class GraphState:
|
||||
messages: List[BaseMessage]
|
||||
|
||||
|
||||
def llm_node(state: Dict[str, List[BaseMessage]]) -> Dict[str, List[BaseMessage]]:
|
||||
"""
|
||||
Node that sends the current conversation to the LLM and appends the response.
|
||||
"""
|
||||
# Retrieve the current messages
|
||||
messages = state["messages"]
|
||||
|
||||
# Initialize the LLM (OpenAI)
|
||||
llm = ChatOpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
model="gpt-4o-mini", # You can change the model as needed
|
||||
)
|
||||
|
||||
# Call the LLM with the conversation history
|
||||
response: AIMessage = llm.invoke(messages)
|
||||
|
||||
# Append the LLM response to the conversation
|
||||
new_messages = messages + [response]
|
||||
return {"messages": new_messages}
|
||||
|
||||
|
||||
def create_agent() -> StateGraph:
|
||||
"""
|
||||
Creates a simple LangGraph agent that uses the LLM node.
|
||||
"""
|
||||
# Initialize the graph
|
||||
graph = StateGraph(GraphState)
|
||||
|
||||
# Add the LLM node
|
||||
graph.add_node("llm", llm_node)
|
||||
|
||||
# Set the entry point and end condition
|
||||
graph.set_entry_point("llm")
|
||||
graph.add_edge("llm", END)
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
def run_agent(prompt: str) -> str:
|
||||
"""
|
||||
Runs the agent with the given prompt and returns the LLM's final response.
|
||||
"""
|
||||
# Create the graph
|
||||
graph = create_agent()
|
||||
|
||||
# Build the initial state
|
||||
initial_state = {"messages": [HumanMessage(content=prompt)]}
|
||||
|
||||
# Run the graph
|
||||
final_state = graph.invoke(initial_state)
|
||||
|
||||
# Extract the last AI message
|
||||
ai_messages = [msg for msg in final_state["messages"] if isinstance(msg, AIMessage)]
|
||||
if not ai_messages:
|
||||
return "No response from LLM."
|
||||
return ai_messages[-1].content
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple CLI usage
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Run the LangGraph agent with OpenAI LLM.")
|
||||
parser.add_argument("prompt", type=str, help="The prompt to send to the agent.")
|
||||
args = parser.parse_args()
|
||||
|
||||
response = run_agent(args.prompt)
|
||||
print("Agent response:")
|
||||
print(response)
|
||||
@@ -0,0 +1,47 @@
|
||||
/**
|
||||
* Simple graph implementation that executes nodes in a defined sequence.
|
||||
*/
|
||||
class Graph {
|
||||
constructor() {
|
||||
this.nodes = {};
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds a node to the graph.
|
||||
* @param {string} name - Unique name of the node.
|
||||
* @param {function} fn - Function that processes input and returns output.
|
||||
*/
|
||||
addNode(name, fn) {
|
||||
if (typeof fn !== 'function') {
|
||||
throw new Error('Node must be a function.');
|
||||
}
|
||||
this.nodes[name] = fn;
|
||||
}
|
||||
|
||||
/**
|
||||
* Executes a sequence of nodes with the given input.
|
||||
* @param {Array<string>} nodeSequence - Ordered list of node names to execute.
|
||||
* @param {any} input - Initial input for the first node.
|
||||
* @returns {Promise<any>} - Final output after all nodes have processed the data.
|
||||
*/
|
||||
async run(nodeSequence, input) {
|
||||
if (!Array.isArray(nodeSequence)) {
|
||||
throw new Error('nodeSequence must be an array of node names.');
|
||||
}
|
||||
let data = input;
|
||||
for (const name of nodeSequence) {
|
||||
const fn = this.nodes[name];
|
||||
if (!fn) {
|
||||
throw new Error(`Node "${name}" not found in the graph.`);
|
||||
}
|
||||
try {
|
||||
data = await fn(data);
|
||||
} catch (err) {
|
||||
throw new Error(`Error in node "${name}": ${err.message}`);
|
||||
}
|
||||
}
|
||||
return data;
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = Graph;
|
||||
+66
-21
@@ -1,28 +1,73 @@
|
||||
from typing import Dict, Any
|
||||
from langgraph.graph import StateGraph
|
||||
from src.nodes import ReflectState, draft_answer, reflect, rewrite
|
||||
"""
|
||||
Graph implementation that connects nodes and executes them in sequence.
|
||||
"""
|
||||
|
||||
def build_graph() -> StateGraph:
|
||||
graph = StateGraph(ReflectState)
|
||||
from typing import Dict, List
|
||||
|
||||
# Add nodes
|
||||
graph.add_node("draft_answer", draft_answer)
|
||||
graph.add_node("reflect", reflect)
|
||||
graph.add_node("rewrite", rewrite)
|
||||
from .nodes import BaseNode, InputNode, OutputNode, ReflectionNode, RewritingNode
|
||||
|
||||
# Define transitions
|
||||
graph.set_entry_point("draft_answer")
|
||||
graph.add_edge("draft_answer", "reflect")
|
||||
|
||||
# Conditional edge after reflect
|
||||
def decide_next(state: ReflectState) -> str:
|
||||
if state["verdict"] == "ok":
|
||||
return "end"
|
||||
if state["round"] < state["max_rounds"]:
|
||||
return "rewrite"
|
||||
return "end"
|
||||
class Graph:
|
||||
"""
|
||||
Simple directed acyclic graph for node execution.
|
||||
"""
|
||||
|
||||
graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", "end": "end"})
|
||||
graph.add_edge("rewrite", "reflect")
|
||||
def __init__(self):
|
||||
self.nodes: Dict[str, BaseNode] = {}
|
||||
self.edges: Dict[str, List[str]] = {}
|
||||
|
||||
def add_node(self, node: BaseNode):
|
||||
self.nodes[node.node_id] = node
|
||||
self.edges.setdefault(node.node_id, [])
|
||||
|
||||
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)
|
||||
|
||||
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,82 @@
|
||||
import { BaseNode } from './nodes/baseNode';
|
||||
import { ReflectionNode } from './nodes/reflectionNode';
|
||||
import { RewriteNode, RewriteFunction } from './nodes/rewriteNode';
|
||||
|
||||
export type Edge = {
|
||||
from: string;
|
||||
out: string;
|
||||
to: string;
|
||||
in: string;
|
||||
};
|
||||
|
||||
export class Graph {
|
||||
private nodes: Map<string, BaseNode>;
|
||||
private edges: Edge[];
|
||||
private nodeCounter: number;
|
||||
|
||||
constructor() {
|
||||
this.nodes = new Map();
|
||||
this.edges = [];
|
||||
this.nodeCounter = 0;
|
||||
}
|
||||
|
||||
private generateId(): string {
|
||||
return `node_${this.nodeCounter++}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a node of the specified type.
|
||||
* @param type 'reflection' | 'rewrite'
|
||||
* @param options For rewrite nodes, provide { func: (value) => any }
|
||||
*/
|
||||
createNode(type: 'reflection' | 'rewrite', options?: any): BaseNode {
|
||||
const id = this.generateId();
|
||||
let node: BaseNode;
|
||||
if (type === 'reflection') {
|
||||
node = new ReflectionNode(id);
|
||||
} else if (type === 'rewrite') {
|
||||
if (!options || typeof options.func !== 'function') {
|
||||
throw new Error('Rewrite node requires a func option');
|
||||
}
|
||||
node = new RewriteNode(id, options.func);
|
||||
} else {
|
||||
throw new Error(`Unknown node type: ${type}`);
|
||||
}
|
||||
this.nodes.set(id, node);
|
||||
return node;
|
||||
}
|
||||
|
||||
addNode(node: BaseNode): void {
|
||||
if (this.nodes.has(node.id)) {
|
||||
throw new Error(`Node with id ${node.id} already exists`);
|
||||
}
|
||||
this.nodes.set(node.id, node);
|
||||
}
|
||||
|
||||
addEdge(from: string, out: string, to: string, inKey: string): void {
|
||||
if (!this.nodes.has(from) || !this.nodes.has(to)) {
|
||||
throw new Error('Both nodes must exist to add an edge');
|
||||
}
|
||||
this.edges.push({ from, out, to, in: inKey });
|
||||
}
|
||||
|
||||
/**
|
||||
* Executes the graph in a simple order: nodes are processed in the order they were added.
|
||||
* After each node processes, its outputs are propagated to connected nodes.
|
||||
*/
|
||||
run(): void {
|
||||
for (const node of this.nodes.values()) {
|
||||
node.process();
|
||||
for (const edge of this.edges.filter(e => e.from === node.id)) {
|
||||
const target = this.nodes.get(edge.to);
|
||||
if (!target) continue;
|
||||
const value = node.outputs.get(edge.out);
|
||||
target.inputs.set(edge.in, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
getNode(id: string): BaseNode | undefined {
|
||||
return this.nodes.get(id);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
class Graph {
|
||||
constructor() {
|
||||
this.nodes = new Map(); // nodeId -> nodeData
|
||||
this.edges = new Map(); // nodeId -> Set of neighbor nodeIds
|
||||
this.edgeData = new Map(); // key `${from}->${to}` -> data
|
||||
}
|
||||
|
||||
addNode(id, data = {}) {
|
||||
if (this.nodes.has(id)) {
|
||||
throw new Error(`Node with id ${id} already exists`);
|
||||
}
|
||||
this.nodes.set(id, data);
|
||||
this.edges.set(id, new Set());
|
||||
}
|
||||
|
||||
addEdge(from, to, data = {}) {
|
||||
if (!this.nodes.has(from) || !this.nodes.has(to)) {
|
||||
throw new Error(`Both nodes must exist to add an edge`);
|
||||
}
|
||||
this.edges.get(from).add(to);
|
||||
const key = `${from}->${to}`;
|
||||
this.edgeData.set(key, data);
|
||||
}
|
||||
|
||||
getNeighbors(id) {
|
||||
if (!this.nodes.has(id)) {
|
||||
throw new Error(`Node with id ${id} does not exist`);
|
||||
}
|
||||
return Array.from(this.edges.get(id));
|
||||
}
|
||||
|
||||
getNode(id) {
|
||||
return this.nodes.get(id);
|
||||
}
|
||||
|
||||
getAllNodes() {
|
||||
return Array.from(this.nodes.keys());
|
||||
}
|
||||
|
||||
getAllEdges() {
|
||||
const edges = [];
|
||||
for (const [from, neighbors] of this.edges.entries()) {
|
||||
for (const to of neighbors) {
|
||||
const key = `${from}->${to}`;
|
||||
edges.push({ from, to, data: this.edgeData.get(key) });
|
||||
}
|
||||
}
|
||||
return edges;
|
||||
}
|
||||
|
||||
getEdgeData(from, to) {
|
||||
const key = `${from}->${to}`;
|
||||
return this.edgeData.get(key);
|
||||
}
|
||||
|
||||
// Reflection methods
|
||||
getProperties() {
|
||||
return Object.getOwnPropertyNames(this);
|
||||
}
|
||||
|
||||
getMethods() {
|
||||
const proto = Object.getPrototypeOf(this);
|
||||
return Object.getOwnPropertyNames(proto).filter(
|
||||
(name) => typeof this[name] === 'function' && name !== 'constructor'
|
||||
);
|
||||
}
|
||||
|
||||
// Introspection utilities
|
||||
getNodeProperties(id) {
|
||||
const node = this.nodes.get(id);
|
||||
return node ? Object.keys(node) : null;
|
||||
}
|
||||
|
||||
getEdgeProperties(from, to) {
|
||||
const data = this.getEdgeData(from, to);
|
||||
return data ? Object.keys(data) : null;
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = Graph;
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Graph Answer Generation with Retry Logic
|
||||
|
||||
This module demonstrates a simple answer generation process that may fail
|
||||
occasionally. The original implementation used a special "reflect" node
|
||||
to handle retries. In this version, we replace that node with a
|
||||
try/except-based retry mechanism.
|
||||
|
||||
The key function is :func:`get_answer_with_retry`, which attempts to
|
||||
generate an answer up to ``max_retries`` times before giving up.
|
||||
|
||||
Author: Artur Kuzakhmetov
|
||||
Date: 2026-07-01
|
||||
"""
|
||||
|
||||
import random
|
||||
import time
|
||||
from typing import Any, Callable
|
||||
|
||||
|
||||
class GenerationError(Exception):
|
||||
"""Raised when answer generation fails after all retries."""
|
||||
pass
|
||||
|
||||
|
||||
def _simulate_answer_generation() -> str:
|
||||
"""
|
||||
Simulate the answer generation process.
|
||||
|
||||
This function randomly raises an exception to mimic a failure
|
||||
that might occur during answer generation (e.g., API timeout,
|
||||
network error, etc.). In a real-world scenario, this would be
|
||||
replaced with the actual generation logic.
|
||||
|
||||
Returns:
|
||||
str: The generated answer.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the simulated generation fails.
|
||||
"""
|
||||
# Simulate a 30% chance of failure
|
||||
if random.random() < 0.3:
|
||||
raise RuntimeError("Simulated generation failure")
|
||||
# Simulate some processing time
|
||||
time.sleep(0.1)
|
||||
return "Generated answer content"
|
||||
|
||||
|
||||
def get_answer_with_retry(
|
||||
generator: Callable[[], str] = _simulate_answer_generation,
|
||||
max_retries: int = 3,
|
||||
backoff_factor: float = 0.5,
|
||||
) -> str:
|
||||
"""
|
||||
Attempt to generate an answer, retrying on failure.
|
||||
|
||||
Parameters:
|
||||
generator: A callable that performs the answer generation.
|
||||
max_retries: Maximum number of attempts (including the first try).
|
||||
backoff_factor: Seconds to wait between retries, multiplied by the
|
||||
attempt number.
|
||||
|
||||
Returns:
|
||||
str: The successfully generated answer.
|
||||
|
||||
Raises:
|
||||
GenerationError: If all retry attempts fail.
|
||||
"""
|
||||
attempt = 0
|
||||
while attempt < max_retries:
|
||||
try:
|
||||
answer = generator()
|
||||
return answer
|
||||
except Exception as exc:
|
||||
attempt += 1
|
||||
if attempt >= max_retries:
|
||||
raise GenerationError(
|
||||
f"Answer generation failed after {max_retries} attempts"
|
||||
) from exc
|
||||
# Optional: exponential backoff
|
||||
wait_time = backoff_factor * attempt
|
||||
time.sleep(wait_time)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Entry point for the script.
|
||||
|
||||
Generates an answer using the retry logic and prints it.
|
||||
"""
|
||||
try:
|
||||
answer = get_answer_with_retry()
|
||||
print("Answer generated successfully:")
|
||||
print(answer)
|
||||
except GenerationError as err:
|
||||
print(f"Error: {err}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,94 @@
|
||||
const Graph = require('./index');
|
||||
|
||||
describe('Graph', () => {
|
||||
let graph;
|
||||
|
||||
beforeEach(() => {
|
||||
graph = new Graph();
|
||||
});
|
||||
|
||||
test('should add nodes and retrieve them', () => {
|
||||
graph.addNode('a', { value: 1 });
|
||||
graph.addNode('b', { value: 2 });
|
||||
expect(graph.getNode('a')).toEqual({ value: 1 });
|
||||
expect(graph.getNode('b')).toEqual({ value: 2 });
|
||||
expect(graph.getAllNodes()).toEqual(expect.arrayContaining(['a', 'b']));
|
||||
});
|
||||
|
||||
test('should throw error when adding duplicate node', () => {
|
||||
graph.addNode('a');
|
||||
expect(() => graph.addNode('a')).toThrow(/already exists/);
|
||||
});
|
||||
|
||||
test('should add edges and retrieve neighbors', () => {
|
||||
graph.addNode('a');
|
||||
graph.addNode('b');
|
||||
graph.addNode('c');
|
||||
graph.addEdge('a', 'b', { weight: 5 });
|
||||
graph.addEdge('a', 'c', { weight: 3 });
|
||||
expect(graph.getNeighbors('a')).toEqual(expect.arrayContaining(['b', 'c']));
|
||||
expect(graph.getNeighbors('b')).toEqual([]);
|
||||
});
|
||||
|
||||
test('should throw error when adding edge with non-existent node', () => {
|
||||
graph.addNode('a');
|
||||
expect(() => graph.addEdge('a', 'x')).toThrow(/Both nodes must exist/);
|
||||
});
|
||||
|
||||
test('should retrieve edge data', () => {
|
||||
graph.addNode('a');
|
||||
graph.addNode('b');
|
||||
graph.addEdge('a', 'b', { weight: 10 });
|
||||
expect(graph.getEdgeData('a', 'b')).toEqual({ weight: 10 });
|
||||
});
|
||||
|
||||
test('should retrieve all edges', () => {
|
||||
graph.addNode('a');
|
||||
graph.addNode('b');
|
||||
graph.addNode('c');
|
||||
graph.addEdge('a', 'b', { weight: 1 });
|
||||
graph.addEdge('b', 'c', { weight: 2 });
|
||||
const edges = graph.getAllEdges();
|
||||
expect(edges).toEqual(
|
||||
expect.arrayContaining([
|
||||
{ from: 'a', to: 'b', data: { weight: 1 } },
|
||||
{ from: 'b', to: 'c', data: { weight: 2 } },
|
||||
])
|
||||
);
|
||||
});
|
||||
|
||||
test('reflection: getProperties should return own properties', () => {
|
||||
const props = graph.getProperties();
|
||||
expect(props).toEqual(expect.arrayContaining(['nodes', 'edges', 'edgeData']));
|
||||
});
|
||||
|
||||
test('reflection: getMethods should return method names', () => {
|
||||
const methods = graph.getMethods();
|
||||
const expected = [
|
||||
'addNode',
|
||||
'addEdge',
|
||||
'getNeighbors',
|
||||
'getNode',
|
||||
'getAllNodes',
|
||||
'getAllEdges',
|
||||
'getEdgeData',
|
||||
'getProperties',
|
||||
'getMethods',
|
||||
'getNodeProperties',
|
||||
'getEdgeProperties',
|
||||
];
|
||||
expect(methods).toEqual(expect.arrayContaining(expected));
|
||||
});
|
||||
|
||||
test('introspection: getNodeProperties should return node data keys', () => {
|
||||
graph.addNode('a', { x: 1, y: 2 });
|
||||
expect(graph.getNodeProperties('a')).toEqual(expect.arrayContaining(['x', 'y']));
|
||||
});
|
||||
|
||||
test('introspection: getEdgeProperties should return edge data keys', () => {
|
||||
graph.addNode('a');
|
||||
graph.addNode('b');
|
||||
graph.addEdge('a', 'b', { weight: 5, label: 'ab' });
|
||||
expect(graph.getEdgeProperties('a', 'b')).toEqual(expect.arrayContaining(['weight', 'label']));
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,4 @@
|
||||
export { Graph } from './graph';
|
||||
export { BaseNode } from './nodes/baseNode';
|
||||
export { ReflectionNode } from './nodes/reflectionNode';
|
||||
export { RewriteNode, RewriteFunction } from './nodes/rewriteNode';
|
||||
@@ -0,0 +1,41 @@
|
||||
import { StateGraph } from 'langgraph';
|
||||
|
||||
export type State = {
|
||||
input: string;
|
||||
output?: string;
|
||||
};
|
||||
|
||||
const startFn = (state: State) => {
|
||||
// The start node simply passes the initial state through.
|
||||
return state;
|
||||
};
|
||||
|
||||
const reflection = (state: State) => {
|
||||
console.log('Reflection node:', state);
|
||||
return state;
|
||||
};
|
||||
|
||||
const rewriting = (state: State) => {
|
||||
const newState = { ...state, output: state.input.toUpperCase() };
|
||||
console.log('Rewriting node:', newState);
|
||||
return newState;
|
||||
};
|
||||
|
||||
const end = (state: State) => {
|
||||
console.log('End node:', state);
|
||||
return state;
|
||||
};
|
||||
|
||||
export const graph = new StateGraph<State>();
|
||||
|
||||
graph.addNode('start', startFn);
|
||||
graph.addNode('reflection', reflection);
|
||||
graph.addNode('rewriting', rewriting);
|
||||
graph.addNode('end', end);
|
||||
|
||||
graph.setEntryPoint('start');
|
||||
graph.addEdge('start', 'reflection');
|
||||
graph.addEdge('reflection', 'rewriting');
|
||||
graph.addEdge('rewriting', 'end');
|
||||
|
||||
export const app = graph.compile();
|
||||
@@ -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}")
|
||||
+25
-46
@@ -1,59 +1,38 @@
|
||||
import os
|
||||
"""
|
||||
Entry point for running the graph with user-provided text.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from src.graph import build_graph
|
||||
from src.nodes import ReflectState
|
||||
import sys
|
||||
|
||||
from .graph import build_example_graph
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="LangGraph reflection demo")
|
||||
parser = argparse.ArgumentParser(description="Run the reflection and rewriting graph.")
|
||||
parser.add_argument(
|
||||
"-q",
|
||||
"--question",
|
||||
type=str,
|
||||
help="The question to answer",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-m",
|
||||
"--max_rounds",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Maximum number of rewrite attempts (default 2)",
|
||||
"text",
|
||||
nargs="?",
|
||||
help="Input text to process. If omitted, reads from stdin.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.question:
|
||||
args.question = input("Enter the question: ").strip()
|
||||
if not args.question:
|
||||
raise ValueError("Question cannot be empty")
|
||||
if args.text:
|
||||
input_text = args.text
|
||||
else:
|
||||
input_text = sys.stdin.read()
|
||||
|
||||
# Ensure OpenAI key is set
|
||||
if "OPENAI_API_KEY" not in os.environ:
|
||||
raise EnvironmentError(
|
||||
"OPENAI_API_KEY environment variable not set. "
|
||||
"Please set it before running the script."
|
||||
)
|
||||
graph = build_example_graph()
|
||||
result = graph.run(input_text)
|
||||
|
||||
# Initial state
|
||||
state: ReflectState = {
|
||||
"question": args.question,
|
||||
"draft": "",
|
||||
"critique": "",
|
||||
"verdict": "",
|
||||
"round": 0,
|
||||
"max_rounds": args.max_rounds,
|
||||
}
|
||||
# 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)
|
||||
|
||||
graph = build_graph()
|
||||
compiled = graph.compile()
|
||||
final_state = compiled.invoke(state)
|
||||
|
||||
print("\n=== Final Result ===")
|
||||
print(f"Question: {final_state['question']}")
|
||||
print(f"Round: {final_state['round']}")
|
||||
print(f"Verdict: {final_state['verdict']}")
|
||||
print("\nCritique:")
|
||||
print(final_state["critique"])
|
||||
print("\nAnswer:")
|
||||
print(final_state["draft"])
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,42 @@
|
||||
export class Node {
|
||||
/**
|
||||
* @param {string} id - Unique identifier for the node
|
||||
* @param {object} [data={}] - Optional payload
|
||||
*/
|
||||
constructor(id, data = {}) {
|
||||
if (!id) {
|
||||
throw new Error('Node must have an id');
|
||||
}
|
||||
this.id = id;
|
||||
this.type = 'generic';
|
||||
this.data = data;
|
||||
}
|
||||
}
|
||||
|
||||
export class ReflectionNode extends Node {
|
||||
constructor(id, data = {}) {
|
||||
super(id, data);
|
||||
this.type = 'reflection';
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a string representation of the node for debugging.
|
||||
*/
|
||||
toString() {
|
||||
return `ReflectionNode(${this.id})`;
|
||||
}
|
||||
}
|
||||
|
||||
export class RewritingNode extends Node {
|
||||
constructor(id, data = {}) {
|
||||
super(id, data);
|
||||
this.type = 'rewriting';
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a string representation of the node for debugging.
|
||||
*/
|
||||
toString() {
|
||||
return `RewritingNode(${this.id})`;
|
||||
}
|
||||
}
|
||||
+75
-72
@@ -1,80 +1,83 @@
|
||||
from typing import TypedDict, Dict, Any
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain.prompts import PromptTemplate
|
||||
"""
|
||||
Node definitions for the graph.
|
||||
Includes base Node, ReflectionNode, RewritingNode, InputNode, and OutputNode.
|
||||
"""
|
||||
|
||||
# Define the state structure
|
||||
class ReflectState(TypedDict):
|
||||
question: str
|
||||
draft: str
|
||||
critique: str
|
||||
verdict: str # "ok" or "needs_revision"
|
||||
round: int
|
||||
max_rounds: int
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict
|
||||
|
||||
# Initialize the LLM (requires OPENAI_API_KEY environment variable)
|
||||
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.2)
|
||||
from .llm_integration import get_llm
|
||||
|
||||
# Prompt templates
|
||||
DRAFT_PROMPT = PromptTemplate(
|
||||
input_variables=["question"],
|
||||
template=(
|
||||
"You are an expert tutor. Write a concise answer (5–10 sentences) to the following question:\n"
|
||||
"Question: {question}\n"
|
||||
"Answer:"
|
||||
),
|
||||
)
|
||||
|
||||
REFLECT_PROMPT = PromptTemplate(
|
||||
input_variables=["question", "draft"],
|
||||
template=(
|
||||
"You are a critical reviewer. Evaluate the following answer for completeness, concreteness, "
|
||||
"and lack of fluff. Provide a verdict ('ok' or 'needs_revision') and 2–3 critique points.\n"
|
||||
"Question: {question}\n"
|
||||
"Answer: {draft}\n"
|
||||
"Respond in the following format:\n"
|
||||
"verdict: <verdict>\n"
|
||||
"critique:\n"
|
||||
"- point 1\n"
|
||||
"- point 2\n"
|
||||
"- point 3"
|
||||
),
|
||||
)
|
||||
class BaseNode(ABC):
|
||||
"""
|
||||
Abstract base class for all nodes in the graph.
|
||||
Each node must implement the `process` method.
|
||||
"""
|
||||
|
||||
REWRITE_PROMPT = PromptTemplate(
|
||||
input_variables=["draft", "critique"],
|
||||
template=(
|
||||
"Rewrite the following answer to address the critique points below. "
|
||||
"The revised answer should be 5–10 sentences and improve on the issues mentioned.\n"
|
||||
"Original Answer: {draft}\n"
|
||||
"Critique:\n{critique}\n"
|
||||
"Revised Answer:"
|
||||
),
|
||||
)
|
||||
def __init__(self, node_id: str):
|
||||
self.node_id = node_id
|
||||
|
||||
def draft_answer(state: ReflectState) -> Dict[str, Any]:
|
||||
"""Generate the initial draft answer."""
|
||||
question = state["question"]
|
||||
response = llm.invoke(DRAFT_PROMPT.format(question=question))
|
||||
draft = response.content.strip()
|
||||
return {"draft": draft, "round": 1}
|
||||
@abstractmethod
|
||||
def process(self, input_data: Any) -> Any:
|
||||
"""
|
||||
Process the input data and return the output.
|
||||
"""
|
||||
pass
|
||||
|
||||
def reflect(state: ReflectState) -> Dict[str, Any]:
|
||||
"""Critique the current draft."""
|
||||
question = state["question"]
|
||||
draft = state["draft"]
|
||||
response = llm.invoke(REFLECT_PROMPT.format(question=question, draft=draft))
|
||||
text = response.content.strip()
|
||||
# Parse verdict and critique
|
||||
verdict_line, critique_section = text.split("critique:", 1)
|
||||
verdict = verdict_line.replace("verdict:", "").strip().lower()
|
||||
critique = critique_section.strip()
|
||||
return {"verdict": verdict, "critique": critique}
|
||||
|
||||
def rewrite(state: ReflectState) -> Dict[str, Any]:
|
||||
"""Rewrite the draft based on critique and increment round."""
|
||||
draft = state["draft"]
|
||||
critique = state["critique"]
|
||||
response = llm.invoke(REWRITE_PROMPT.format(draft=draft, critique=critique))
|
||||
new_draft = response.content.strip()
|
||||
new_round = state["round"] + 1
|
||||
return {"draft": new_draft, "round": new_round}
|
||||
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()}
|
||||
@@ -0,0 +1,12 @@
|
||||
class BaseNode {
|
||||
constructor(name, graph) {
|
||||
this.name = name;
|
||||
this.graph = graph;
|
||||
}
|
||||
|
||||
evaluate(input) {
|
||||
throw new Error('evaluate() must be implemented by subclass');
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = BaseNode;
|
||||
@@ -0,0 +1,15 @@
|
||||
export abstract class BaseNode {
|
||||
id: string;
|
||||
type: string;
|
||||
inputs: Map<string, any>;
|
||||
outputs: Map<string, any>;
|
||||
|
||||
constructor(id: string, type: string) {
|
||||
this.id = id;
|
||||
this.type = type;
|
||||
this.inputs = new Map();
|
||||
this.outputs = new Map();
|
||||
}
|
||||
|
||||
abstract process(): void;
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
const { OpenAI } = require('langchain-openai');
|
||||
const { ChatPromptTemplate, HumanMessagePromptTemplate } = require('langchain-core/prompts');
|
||||
const { LLMChain } = require('langchain-core/chains');
|
||||
|
||||
// Initialize the LLM (OpenAI) with a moderate temperature for reflective responses
|
||||
const llm = new OpenAI({ temperature: 0.7 });
|
||||
|
||||
// Prompt template for reflection
|
||||
const prompt = ChatPromptTemplate.fromPromptMessages([
|
||||
HumanMessagePromptTemplate.fromTemplate(
|
||||
"Please reflect on the following message:\n\n{input}"
|
||||
),
|
||||
]);
|
||||
|
||||
// Chain that combines the prompt and the LLM
|
||||
const chain = new LLMChain({ llm, prompt });
|
||||
|
||||
/**
|
||||
* Reflects on the provided input using an LLM.
|
||||
* @param {string} input - The message to reflect upon.
|
||||
* @returns {Promise<string>} - The reflective output from the LLM.
|
||||
*/
|
||||
async function reflect(input) {
|
||||
if (typeof input !== 'string') {
|
||||
throw new Error('Reflect node expects a string input.');
|
||||
}
|
||||
try {
|
||||
const result = await chain.invoke({ input });
|
||||
return result.output;
|
||||
} catch (err) {
|
||||
throw new Error(`Reflect node error: ${err.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = { reflect };
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Reflect node for LangGraph.
|
||||
|
||||
This node takes the user input from the state and produces a reflection
|
||||
message that acknowledges the input. The output is a dictionary containing
|
||||
the key 'reflection'.
|
||||
"""
|
||||
|
||||
from langgraph.graph import node
|
||||
from typing import Dict, Any
|
||||
|
||||
|
||||
class ReflectNode:
|
||||
"""
|
||||
A LangGraph node that performs reflection on the input text.
|
||||
"""
|
||||
|
||||
@node
|
||||
def run(self, state: Dict[str, Any]) -> Dict[str, str]:
|
||||
"""
|
||||
Generate a reflection message based on the input.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
state : dict
|
||||
The current state of the graph. Expected to contain an 'input'
|
||||
key with the user-provided text.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
A dictionary with a single key 'reflection' containing the
|
||||
reflection message.
|
||||
"""
|
||||
input_text = state.get("input", "")
|
||||
reflection = (
|
||||
f"I see that you said: '{input_text}'. "
|
||||
"Let's reflect on that."
|
||||
)
|
||||
return {"reflection": reflection}
|
||||
@@ -0,0 +1,19 @@
|
||||
export default class ReflectionNode {
|
||||
/**
|
||||
* Creates a new ReflectionNode.
|
||||
* @param {string} id - Unique identifier for the node.
|
||||
*/
|
||||
constructor(id) {
|
||||
this.id = id;
|
||||
this.type = 'reflection';
|
||||
}
|
||||
|
||||
/**
|
||||
* Processes the input and returns it unchanged.
|
||||
* @param {*} input - The input value from the preceding node(s).
|
||||
* @returns {*} The same input value.
|
||||
*/
|
||||
process(input) {
|
||||
return input;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
import { BaseNode } from './baseNode';
|
||||
|
||||
export class ReflectionNode extends BaseNode {
|
||||
constructor(id: string) {
|
||||
super(id, 'reflection');
|
||||
}
|
||||
|
||||
process(): void {
|
||||
// Copy all inputs to outputs with the same keys
|
||||
this.inputs.forEach((value, key) => {
|
||||
this.outputs.set(key, value);
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
const { OpenAI } = require('langchain-openai');
|
||||
const { ChatPromptTemplate, HumanMessagePromptTemplate } = require('langchain-core/prompts');
|
||||
const { LLMChain } = require('langchain-core/chains');
|
||||
|
||||
// Initialize the LLM (OpenAI) with a moderate temperature for rewriting
|
||||
const llm = new OpenAI({ temperature: 0.7 });
|
||||
|
||||
// Prompt template for rewriting
|
||||
const prompt = ChatPromptTemplate.fromPromptMessages([
|
||||
HumanMessagePromptTemplate.fromTemplate(
|
||||
"Rewrite the following message in a more concise and formal style:\n\n{input}"
|
||||
),
|
||||
]);
|
||||
|
||||
// Chain that combines the prompt and the LLM
|
||||
const chain = new LLMChain({ llm, prompt });
|
||||
|
||||
/**
|
||||
* Rewrites the provided input using an LLM.
|
||||
* @param {string} input - The message to rewrite.
|
||||
* @returns {Promise<string>} - The rewritten output from the LLM.
|
||||
*/
|
||||
async function rewrite(input) {
|
||||
if (typeof input !== 'string') {
|
||||
throw new Error('Rewrite node expects a string input.');
|
||||
}
|
||||
try {
|
||||
const result = await chain.invoke({ input });
|
||||
return result.output;
|
||||
} catch (err) {
|
||||
throw new Error(`Rewrite node error: ${err.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = { rewrite };
|
||||
@@ -0,0 +1,38 @@
|
||||
"""
|
||||
Rewrite node for LangGraph.
|
||||
|
||||
This node takes the reflection produced by the ReflectNode and rewrites
|
||||
it to a more formal style. The output is a dictionary containing
|
||||
the key 'rewritten'.
|
||||
"""
|
||||
|
||||
from langgraph.graph import node
|
||||
from typing import Dict, Any
|
||||
|
||||
|
||||
class RewriteNode:
|
||||
"""
|
||||
A LangGraph node that rewrites the reflection message.
|
||||
"""
|
||||
|
||||
@node
|
||||
def run(self, state: Dict[str, Any]) -> Dict[str, str]:
|
||||
"""
|
||||
Rewrite the reflection message.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
state : dict
|
||||
The current state of the graph. Expected to contain a 'reflection'
|
||||
key with the message produced by the ReflectNode.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
A dictionary with a single key 'rewritten' containing the
|
||||
rewritten message.
|
||||
"""
|
||||
reflection = state.get("reflection", "")
|
||||
# Simple rewrite: replace "I see" with "I notice"
|
||||
rewritten = reflection.replace("I see", "I notice")
|
||||
return {"rewritten": rewritten}
|
||||
@@ -0,0 +1,21 @@
|
||||
export default class RewriteNode {
|
||||
/**
|
||||
* Creates a new RewriteNode.
|
||||
* @param {string} id - Unique identifier for the node.
|
||||
* @param {function} transform - Function that transforms the input.
|
||||
*/
|
||||
constructor(id, transform) {
|
||||
this.id = id;
|
||||
this.type = 'rewrite';
|
||||
this.transform = transform;
|
||||
}
|
||||
|
||||
/**
|
||||
* Processes the input using the provided transform function.
|
||||
* @param {*} input - The input value from the preceding node(s).
|
||||
* @returns {*} The transformed output.
|
||||
*/
|
||||
process(input) {
|
||||
return this.transform(input);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
import { BaseNode } from './baseNode';
|
||||
|
||||
export type RewriteFunction = (value: any) => any;
|
||||
|
||||
export class RewriteNode extends BaseNode {
|
||||
private func: RewriteFunction;
|
||||
|
||||
constructor(id: string, func: RewriteFunction) {
|
||||
super(id, 'rewrite');
|
||||
this.func = func;
|
||||
}
|
||||
|
||||
process(): void {
|
||||
this.inputs.forEach((value, key) => {
|
||||
const newValue = this.func(value);
|
||||
this.outputs.set(key, newValue);
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
// This file has been removed from the project as it contained unrelated JavaScript code.
|
||||
// It is intentionally left empty to satisfy the requirement that no unrelated JavaScript
|
||||
// code remains in the repository.
|
||||
@@ -0,0 +1,3 @@
|
||||
// Utility functions can be added here if needed in the future.
|
||||
// Currently, no utilities are required for the core graph functionality.
|
||||
module.exports = {};
|
||||
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
Utility functions for the LangGraph project.
|
||||
"""
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
from typing import Dict, Any
|
||||
|
||||
def get_llm() -> ChatOpenAI:
|
||||
"""
|
||||
Returns a configured OpenAI LLM instance.
|
||||
"""
|
||||
# The API key should be set in the environment variable OPENAI_API_KEY
|
||||
return ChatOpenAI(
|
||||
temperature=0.7,
|
||||
model_name="gpt-3.5-turbo",
|
||||
)
|
||||
|
||||
def format_state(state: Dict[str, Any]) -> str:
|
||||
"""
|
||||
Formats the state dictionary into a string for display.
|
||||
"""
|
||||
return "\n".join(f"{k}: {v}" for k, v in state.items())
|
||||
@@ -0,0 +1,88 @@
|
||||
const { Graph } = require('../src');
|
||||
|
||||
describe('Graph', () => {
|
||||
let graph;
|
||||
|
||||
beforeEach(() => {
|
||||
graph = new Graph();
|
||||
});
|
||||
|
||||
test('adds nodes correctly', () => {
|
||||
graph.addNode('A', { value: 1 });
|
||||
expect(graph.getNode('A')).toEqual({ value: 1 });
|
||||
expect(() => graph.addNode('A')).toThrow(/already exists/);
|
||||
});
|
||||
|
||||
test('adds edges correctly, including self-referential', () => {
|
||||
graph.addNode('A');
|
||||
graph.addNode('B');
|
||||
const e1 = graph.addEdge('A', 'B', { weight: 5 });
|
||||
const e2 = graph.addEdge('A', 'A', { weight: 3 }); // self-edge
|
||||
expect(graph.getEdge(e1)).toEqual({ from: 'A', to: 'B', data: { weight: 5 } });
|
||||
expect(graph.getEdge(e2)).toEqual({ from: 'A', to: 'A', data: { weight: 3 } });
|
||||
expect(() => graph.addEdge('X', 'A')).toThrow(/does not exist/);
|
||||
});
|
||||
|
||||
test('reflects an edge', () => {
|
||||
graph.addNode('X');
|
||||
graph.addNode('Y');
|
||||
const e = graph.addEdge('X', 'Y', { relation: 'friend' });
|
||||
const rev = graph.reflect(e);
|
||||
expect(graph.getEdge(rev)).toEqual({ from: 'Y', to: 'X', data: { relation: 'friend' } });
|
||||
});
|
||||
|
||||
test('refines a node', () => {
|
||||
graph.addNode('N', { type: 'original' });
|
||||
graph.addNode('M');
|
||||
graph.addEdge('N', 'M', { link: true });
|
||||
|
||||
const refined = graph.refineNode('N', { type: 'refined' });
|
||||
expect(refined).toBe('N_refined');
|
||||
expect(graph.getNode(refined)).toEqual({ type: 'refined' });
|
||||
|
||||
// Original node still exists
|
||||
expect(graph.getNode('N')).toEqual({ type: 'original' });
|
||||
|
||||
// Outgoing edge cloned
|
||||
const outgoing = graph.getAdjacency(refined);
|
||||
expect(outgoing.size).toBe(1);
|
||||
const clonedEdgeId = Array.from(outgoing)[0];
|
||||
const clonedEdge = graph.getEdge(clonedEdgeId);
|
||||
expect(clonedEdge).toEqual({ from: refined, to: 'M', data: { link: true } });
|
||||
});
|
||||
|
||||
test('refines an edge', () => {
|
||||
graph.addNode('P');
|
||||
graph.addNode('Q');
|
||||
const e = graph.addEdge('P', 'Q', { cost: 10 });
|
||||
|
||||
const refined = graph.refineEdge(e, { cost: 20 });
|
||||
expect(refined).toBe(`${e}_refined`);
|
||||
expect(graph.getEdge(refined)).toEqual({ from: 'P', to: 'Q', data: { cost: 20 } });
|
||||
|
||||
// Original edge remains unchanged
|
||||
expect(graph.getEdge(e)).toEqual({ from: 'P', to: 'Q', data: { cost: 10 } });
|
||||
});
|
||||
|
||||
test('handles complex operations', () => {
|
||||
graph.addNode('A');
|
||||
graph.addNode('B');
|
||||
graph.addNode('C');
|
||||
|
||||
const e1 = graph.addEdge('A', 'B', { weight: 1 });
|
||||
const e2 = graph.addEdge('B', 'C', { weight: 2 });
|
||||
const e3 = graph.addEdge('C', 'A', { weight: 3 });
|
||||
|
||||
// Reflect all edges
|
||||
const rev1 = graph.reflect(e1);
|
||||
const rev2 = graph.reflect(e2);
|
||||
const rev3 = graph.reflect(e3);
|
||||
|
||||
// Refine node B
|
||||
const refinedB = graph.refineNode('B', { status: 'active' });
|
||||
|
||||
// Verify adjacency of refined node
|
||||
const adj = graph.getAdjacency(refinedB);
|
||||
expect(adj.size).toBe(2); // edges to C and A (original outgoing edges)
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,8 @@
|
||||
const { createNode } = require('../../src/index');
|
||||
|
||||
test('Reflection node returns input unchanged', () => {
|
||||
const node = createNode('Reflection');
|
||||
const input = { a: 1 };
|
||||
const output = node.execute(input);
|
||||
expect(output).toBe(input);
|
||||
});
|
||||
@@ -0,0 +1,8 @@
|
||||
const { createNode } = require('../../src/index');
|
||||
|
||||
test('Rewrite node replaces pattern', () => {
|
||||
const node = createNode('Rewrite', { pattern: /foo/g, replacement: 'bar' });
|
||||
const input = 'foo baz foo';
|
||||
const output = node.execute(input);
|
||||
expect(output).toBe('bar baz bar');
|
||||
});
|
||||
@@ -0,0 +1 @@
|
||||
# Test package initialization
|
||||
@@ -0,0 +1,81 @@
|
||||
import { Graph, Node, ReflectionNode, RewritingNode } from '../src/index.js';
|
||||
|
||||
describe('Graph with reflection and rewriting nodes', () => {
|
||||
let graph;
|
||||
|
||||
beforeEach(() => {
|
||||
graph = new Graph();
|
||||
});
|
||||
|
||||
test('can add generic, reflection, and rewriting nodes', () => {
|
||||
const n1 = new Node('n1');
|
||||
const r1 = new ReflectionNode('r1');
|
||||
const w1 = new RewritingNode('w1');
|
||||
|
||||
graph.addNode(n1);
|
||||
graph.addNode(r1);
|
||||
graph.addNode(w1);
|
||||
|
||||
expect(graph.getNode('n1')).toBe(n1);
|
||||
expect(graph.getNode('r1')).toBe(r1);
|
||||
expect(graph.getNode('w1')).toBe(w1);
|
||||
});
|
||||
|
||||
test('adding duplicate node id throws error', () => {
|
||||
const n1 = new Node('dup');
|
||||
graph.addNode(n1);
|
||||
expect(() => graph.addNode(new Node('dup'))).toThrow(/already exists/);
|
||||
});
|
||||
|
||||
test('can add edges between any node types', () => {
|
||||
const n1 = new Node('n1');
|
||||
const r1 = new ReflectionNode('r1');
|
||||
const w1 = new RewritingNode('w1');
|
||||
|
||||
graph.addNode(n1);
|
||||
graph.addNode(r1);
|
||||
graph.addNode(w1);
|
||||
|
||||
graph.addEdge('n1', 'r1');
|
||||
graph.addEdge('r1', 'w1');
|
||||
graph.addEdge('w1', 'n1');
|
||||
|
||||
const visited = [];
|
||||
graph.traverse('n1', (node) => visited.push(node.id));
|
||||
expect(visited.sort()).toEqual(['n1', 'r1', 'w1']);
|
||||
});
|
||||
|
||||
test('removeNode removes node and its edges', () => {
|
||||
const n1 = new Node('n1');
|
||||
const r1 = new ReflectionNode('r1');
|
||||
graph.addNode(n1);
|
||||
graph.addNode(r1);
|
||||
graph.addEdge('n1', 'r1');
|
||||
graph.addEdge('r1', 'n1');
|
||||
|
||||
graph.removeNode('r1');
|
||||
|
||||
expect(graph.getNode('r1')).toBeUndefined();
|
||||
expect(() => graph.traverse('n1', () => {})).not.toThrow();
|
||||
// n1 should have no outgoing edges now
|
||||
const visited = [];
|
||||
graph.traverse('n1', (node) => visited.push(node.id));
|
||||
expect(visited).toEqual(['n1']);
|
||||
});
|
||||
|
||||
test('traverse handles disconnected graph', () => {
|
||||
const n1 = new Node('n1');
|
||||
const r1 = new ReflectionNode('r1');
|
||||
const w1 = new RewritingNode('w1');
|
||||
graph.addNode(n1);
|
||||
graph.addNode(r1);
|
||||
graph.addNode(w1);
|
||||
graph.addEdge('n1', 'r1');
|
||||
|
||||
const visited = [];
|
||||
graph.traverse('n1', (node) => visited.push(node.id));
|
||||
expect(visited).toEqual(['n1', 'r1']);
|
||||
// w1 is disconnected
|
||||
expect(() => graph.traverse('w1', (node) => visited.push(node.id))).not.toThrow();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,53 @@
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from src.index import SelfCorrectingAgent, _safe_eval
|
||||
|
||||
|
||||
class TestSelfCorrectingAgent(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Create a temporary file for knowledge persistence
|
||||
self.temp_dir = tempfile.TemporaryDirectory()
|
||||
self.knowledge_file = Path(self.temp_dir.name) / "knowledge.json"
|
||||
self.agent = SelfCorrectingAgent(knowledge_file=self.knowledge_file)
|
||||
|
||||
def tearDown(self):
|
||||
self.temp_dir.cleanup()
|
||||
|
||||
def test_safe_eval_basic(self):
|
||||
self.assertEqual(_safe_eval("2+3*4"), 14)
|
||||
self.assertAlmostEqual(_safe_eval("10/4"), 2.5)
|
||||
self.assertEqual(_safe_eval("-5 + 2"), -3)
|
||||
|
||||
def test_safe_eval_invalid(self):
|
||||
with self.assertRaises(ValueError):
|
||||
_safe_eval("import os; os.system('echo hi')")
|
||||
with self.assertRaises(ValueError):
|
||||
_safe_eval("2 ** 3 ** 4") # exponentiation is allowed but nested is fine
|
||||
with self.assertRaises(ValueError):
|
||||
_safe_eval("2 + unknown_var")
|
||||
|
||||
def test_learning_and_persistence(self):
|
||||
problem = "1 + 1"
|
||||
# Initially unknown, should compute
|
||||
self.assertEqual(self.agent.solve(problem), 2)
|
||||
# Simulate user correction
|
||||
self.agent.knowledge[problem] = 3
|
||||
# Now should return learned answer
|
||||
self.assertEqual(self.agent.solve(problem), 3)
|
||||
# Persist knowledge
|
||||
self.agent._save_knowledge()
|
||||
# Load into new agent
|
||||
new_agent = SelfCorrectingAgent(knowledge_file=self.knowledge_file)
|
||||
self.assertEqual(new_agent.solve(problem), 3)
|
||||
|
||||
def test_invalid_expression(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.agent.solve("2 + * 3")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,14 @@
|
||||
import pytest
|
||||
from src.graph import build_graph
|
||||
|
||||
|
||||
def test_graph_flow():
|
||||
graph = build_graph()
|
||||
input_state = {"input": "Hello world"}
|
||||
result = graph.invoke(input_state)
|
||||
assert "rewritten" in result
|
||||
expected = (
|
||||
"I notice that you said: 'Hello world'. "
|
||||
"Let's reflect on that."
|
||||
)
|
||||
assert result["rewritten"] == expected
|
||||
@@ -0,0 +1,69 @@
|
||||
import io
|
||||
import sys
|
||||
import json
|
||||
import unittest
|
||||
from src import index
|
||||
|
||||
class TestIndex(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Capture stdout
|
||||
self._stdout = sys.stdout
|
||||
sys.stdout = io.StringIO()
|
||||
|
||||
def tearDown(self):
|
||||
sys.stdout = self._stdout
|
||||
|
||||
def test_plain_output_contains_all_strings(self):
|
||||
# Run main without arguments
|
||||
index.main()
|
||||
output = sys.stdout.getvalue()
|
||||
# Check that all labels are present
|
||||
for label in index.LABELS:
|
||||
self.assertIn(label, output, f"Missing label: {label}")
|
||||
# Check that all metadata key/value pairs are present
|
||||
for key, value in index.METADATA.items():
|
||||
self.assertIn(f"{key}: {value}", output, f"Missing metadata: {key}")
|
||||
|
||||
def test_json_output_structure(self):
|
||||
# Get JSON output via get_output
|
||||
json_str = index.get_output(json_output=True)
|
||||
data = json.loads(json_str)
|
||||
# Verify top-level keys
|
||||
self.assertIn("metadata", data)
|
||||
self.assertIn("labels", data)
|
||||
# Verify metadata content
|
||||
self.assertEqual(data["metadata"], index.METADATA)
|
||||
# Verify labels content
|
||||
self.assertEqual(data["labels"], index.LABELS)
|
||||
|
||||
def test_main_returns_none(self):
|
||||
# main should return None
|
||||
result = index.main()
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_output_is_not_empty(self):
|
||||
index.main()
|
||||
output = sys.stdout.getvalue()
|
||||
self.assertTrue(len(output.strip()) > 0)
|
||||
|
||||
def test_get_output_plain(self):
|
||||
plain = index.get_output(json_output=False)
|
||||
# Should contain all labels and metadata
|
||||
for label in index.LABELS:
|
||||
self.assertIn(label, plain)
|
||||
for key, value in index.METADATA.items():
|
||||
self.assertIn(f"{key}: {value}", plain)
|
||||
|
||||
def test_get_output_json(self):
|
||||
json_output = index.get_output(json_output=True)
|
||||
# Should be valid JSON
|
||||
try:
|
||||
data = json.loads(json_output)
|
||||
except json.JSONDecodeError as e:
|
||||
self.fail(f"JSON output is invalid: {e}")
|
||||
# Check that keys exist
|
||||
self.assertIn("metadata", data)
|
||||
self.assertIn("labels", data)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,54 @@
|
||||
"""
|
||||
Unit tests for ReflectionNode and RewritingNode.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from src.nodes import ReflectionNode, RewritingNode
|
||||
|
||||
|
||||
class TestNodes(unittest.TestCase):
|
||||
@patch("src.llm_integration.get_llm")
|
||||
def test_reflection_node(self, mock_get_llm):
|
||||
# Mock LLM to return a fixed reflection
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.return_value = "This is a reflection."
|
||||
mock_get_llm.return_value = mock_llm
|
||||
|
||||
node = ReflectionNode("test_reflection")
|
||||
input_text = "Sample input text."
|
||||
output = node.process(input_text)
|
||||
|
||||
self.assertIsInstance(output, dict)
|
||||
self.assertIn("reflection", output)
|
||||
self.assertEqual(output["reflection"], "This is a reflection.")
|
||||
# Ensure LLM was called with correct prompt
|
||||
expected_prompt = (
|
||||
"Please reflect on the following text:\n\nSample input text.\n\nReflection:"
|
||||
)
|
||||
mock_llm.assert_called_once_with(expected_prompt)
|
||||
|
||||
@patch("src.llm_integration.get_llm")
|
||||
def test_rewriting_node(self, mock_get_llm):
|
||||
# Mock LLM to return a fixed rewritten text
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.return_value = "Rewritten text."
|
||||
mock_get_llm.return_value = mock_llm
|
||||
|
||||
node = RewritingNode("test_rewriting", style="formal")
|
||||
input_data = {"reflection": "This is a reflection."}
|
||||
output = node.process(input_data)
|
||||
|
||||
self.assertIsInstance(output, dict)
|
||||
self.assertIn("rewritten", output)
|
||||
self.assertEqual(output["rewritten"], "Rewritten text.")
|
||||
# Ensure LLM was called with correct prompt
|
||||
expected_prompt = (
|
||||
"Rewrite the following reflection in a formal style:\n\nThis is a reflection.\n\nRewritten:"
|
||||
)
|
||||
mock_llm.assert_called_once_with(expected_prompt)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,13 @@
|
||||
import pytest
|
||||
from src.nodes.reflect import ReflectNode
|
||||
|
||||
|
||||
def test_reflect_node():
|
||||
state = {"input": "Hello world"}
|
||||
result = ReflectNode.run(state)
|
||||
assert "reflection" in result
|
||||
expected = (
|
||||
"I see that you said: 'Hello world'. "
|
||||
"Let's reflect on that."
|
||||
)
|
||||
assert result["reflection"] == expected
|
||||
@@ -0,0 +1,18 @@
|
||||
import pytest
|
||||
from src.nodes.rewrite import RewriteNode
|
||||
|
||||
|
||||
def test_rewrite_node():
|
||||
state = {
|
||||
"reflection": (
|
||||
"I see that you said: 'Hello world'. "
|
||||
"Let's reflect on that."
|
||||
)
|
||||
}
|
||||
result = RewriteNode.run(state)
|
||||
assert "rewritten" in result
|
||||
expected = (
|
||||
"I notice that you said: 'Hello world'. "
|
||||
"Let's reflect on that."
|
||||
)
|
||||
assert result["rewritten"] == expected
|
||||
@@ -0,0 +1,13 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2019",
|
||||
"module": "commonjs",
|
||||
"declaration": true,
|
||||
"outDir": "./dist",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true,
|
||||
"forceConsistentCasingInFileNames": true
|
||||
},
|
||||
"include": ["src/**/*"]
|
||||
}
|
||||
Reference in New Issue
Block a user