56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
from typing import TypedDict, Dict, Any
|
||
from langchain_community.llms import OpenAI
|
||
import re
|
||
|
||
# Global LLM instance
|
||
llm = OpenAI(temperature=0.7)
|
||
|
||
class ReflectState(TypedDict):
|
||
question: str
|
||
draft: str
|
||
critique: str
|
||
verdict: str
|
||
round: int
|
||
max_rounds: int
|
||
|
||
def draft_answer(state: ReflectState) -> ReflectState:
|
||
prompt = (
|
||
f"Answer the following question in 5–10 sentences:\n\n"
|
||
f"Question: {state['question']}\n\n"
|
||
f"Answer:"
|
||
)
|
||
answer = llm.invoke(prompt).strip()
|
||
state["draft"] = answer
|
||
return state
|
||
|
||
def reflect(state: ReflectState) -> ReflectState:
|
||
prompt = (
|
||
f"You are a critical reviewer. Evaluate the following draft answer.\n\n"
|
||
f"Draft:\n{state['draft']}\n\n"
|
||
f"Provide a verdict ('ok' or 'needs_revision') and 2–3 remarks.\n"
|
||
f"Format:\n"
|
||
f"Verdict: <ok|needs_revision>\n"
|
||
f"Remarks:\n<remarks>"
|
||
)
|
||
response = llm.invoke(prompt).strip()
|
||
# Parse verdict
|
||
verdict_match = re.search(r"Verdict:\s*(\w+)", response, re.IGNORECASE)
|
||
remarks_match = re.search(r"Remarks:\s*(.*)", response, re.DOTALL | re.IGNORECASE)
|
||
verdict = verdict_match.group(1).lower() if verdict_match else "needs_revision"
|
||
remarks = remarks_match.group(1).strip() if remarks_match else "No remarks provided."
|
||
state["verdict"] = verdict
|
||
state["critique"] = remarks
|
||
return state
|
||
|
||
def rewrite(state: ReflectState) -> ReflectState:
|
||
prompt = (
|
||
f"Rewrite the following draft answer to address the remarks below.\n\n"
|
||
f"Draft:\n{state['draft']}\n\n"
|
||
f"Remarks:\n{state['critique']}\n\n"
|
||
f"Provide the revised answer in 5–10 sentences."
|
||
)
|
||
revised = llm.invoke(prompt).strip()
|
||
state["draft"] = revised
|
||
state["round"] += 1
|
||
return state
|