diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..ab1f72c --- /dev/null +++ b/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']}")