128 lines
3.8 KiB
Python
128 lines
3.8 KiB
Python
import os
|
|
from typing import TypedDict, Dict, Any
|
|
import json
|
|
|
|
from langgraph import StateGraph, END
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
# State definition
|
|
class ReflectState(TypedDict):
|
|
question: str
|
|
draft: str
|
|
critique: str
|
|
verdict: str # "ok" | "needs_revision"
|
|
round: int
|
|
max_rounds: int
|
|
|
|
# LLM configuration
|
|
|
|
def get_llm() -> ChatOpenAI:
|
|
"""Return a ChatOpenAI instance configured via environment variables.
|
|
If OLLAMA_BASE_URL is set, use Ollama; otherwise try OPENAI_API_KEY; fallback to local.
|
|
"""
|
|
ollama_url = os.getenv("OLLAMA_BASE_URL")
|
|
openai_key = os.getenv("OPENAI_API_KEY")
|
|
if ollama_url:
|
|
return ChatOpenAI(model="llama3", base_url=ollama_url, api_key="ollama")
|
|
if openai_key:
|
|
return ChatOpenAI(model="gpt-4o-mini")
|
|
return ChatOpenAI(model="gpt-4o-mini")
|
|
|
|
llm = get_llm()
|
|
|
|
# Node: draft_answer
|
|
|
|
def draft_answer(state: ReflectState) -> Dict[str, Any]:
|
|
question = state["question"]
|
|
prompt = f"Answer the following question in 5-10 sentences: {question}"
|
|
try:
|
|
response = llm.invoke(prompt)
|
|
draft = str(response)
|
|
except Exception as e:
|
|
raise RuntimeError(f"LLM draft generation failed: {e}")
|
|
return {"draft": draft, "round": 1}
|
|
|
|
# Node: reflect
|
|
|
|
def reflect(state: ReflectState) -> Dict[str, Any]:
|
|
draft = state["draft"]
|
|
prompt = (
|
|
"Critique the following draft for completeness, concreteness, and absence of filler. "
|
|
"Respond with a JSON object containing 'verdict' (values: 'ok' or 'needs_revision') "
|
|
"and 'critique' (text). Example: {\"verdict\": \"ok\", \"critique\": \"...\"}"
|
|
)
|
|
content = f"Draft: {draft}"
|
|
full_prompt = prompt + "\n" + content
|
|
try:
|
|
response = llm.invoke(full_prompt)
|
|
text = str(response)
|
|
except Exception as e:
|
|
raise RuntimeError(f"LLM reflect failed: {e}")
|
|
# Parse JSON
|
|
try:
|
|
data = json.loads(text)
|
|
verdict = data["verdict"]
|
|
critique = data["critique"]
|
|
except Exception:
|
|
# Fallback parsing: treat entire response as critique
|
|
verdict = "needs_revision"
|
|
critique = text
|
|
return {"critique": critique, "verdict": verdict}
|
|
|
|
# Node: rewrite
|
|
|
|
def rewrite(state: ReflectState) -> Dict[str, Any]:
|
|
draft = state["draft"]
|
|
critique = state["critique"]
|
|
prompt = (
|
|
"Rewrite the following draft based on the critique. "
|
|
"Keep the meaning but improve clarity and remove filler. "
|
|
"Output only the revised draft."
|
|
)
|
|
content = f"Draft: {draft}\nCritique: {critique}"
|
|
full_prompt = prompt + "\n" + content
|
|
try:
|
|
response = llm.invoke(full_prompt)
|
|
new_draft = str(response)
|
|
except Exception as e:
|
|
raise RuntimeError(f"LLM rewrite failed: {e}")
|
|
return {"draft": new_draft, "round": state["round"] + 1}
|
|
|
|
# Build graph
|
|
builder = StateGraph(ReflectState)
|
|
|
|
builder.add_node("draft_answer", draft_answer)
|
|
builder.add_node("reflect", reflect)
|
|
builder.add_node("rewrite", rewrite)
|
|
|
|
builder.set_entry_point("draft_answer")
|
|
|
|
# Transitions
|
|
builder.add_edge("draft_answer", "reflect")
|
|
builder.add_conditional_edges(
|
|
"reflect",
|
|
lambda s: END if s["verdict"] == "ok" else "rewrite" if s["round"] < s["max_rounds"] else END,
|
|
)
|
|
builder.add_edge("rewrite", "reflect")
|
|
|
|
# Compile graph
|
|
graph = builder.compile(checkpointer=MemorySaver())
|
|
|
|
# CLI
|
|
if __name__ == "__main__":
|
|
question = input("Enter your question: ")
|
|
initial_state: ReflectState = {
|
|
"question": question,
|
|
"draft": "",
|
|
"critique": "",
|
|
"verdict": "",
|
|
"round": 0,
|
|
"max_rounds": 2,
|
|
}
|
|
result = graph.invoke(initial_state)
|
|
print("\nFinal Result:\n")
|
|
print(f"Draft: {result['draft']}")
|
|
print(f"Critique: {result['critique']}")
|
|
print(f"Verdict: {result['verdict']}")
|