42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import argparse
|
|
import os
|
|
from dotenv import load_dotenv
|
|
|
|
from graph import build_graph
|
|
from nodes import llm
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="LangGraph Reflection Demo")
|
|
parser.add_argument("question", type=str, help="The question to answer")
|
|
parser.add_argument("--max-rounds", type=int, default=2, help="Maximum number of rewrite attempts")
|
|
parser.add_argument("--model", type=str, default="gpt-3.5-turbo", help="LLM model name")
|
|
return parser.parse_args()
|
|
|
|
def main():
|
|
load_dotenv()
|
|
args = parse_args()
|
|
|
|
# Configure LLM
|
|
llm.model = args.model
|
|
|
|
graph = build_graph()
|
|
initial_state = {
|
|
"question": args.question,
|
|
"draft": "",
|
|
"critique": "",
|
|
"verdict": "",
|
|
"round": 1,
|
|
"max_rounds": args.max_rounds,
|
|
}
|
|
|
|
result = graph.invoke(initial_state)
|
|
|
|
print("\n=== Final Result ===")
|
|
print(f"Draft:\n{result['draft']}\n")
|
|
print(f"Verdict: {result['verdict']}")
|
|
print(f"Critique:\n{result['critique']}\n")
|
|
print(f"Round: {result['round']}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|