Add agent.py

agent.py

@@ -0,0 +1,127 @@
+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']}")