Add agent.py

agent.py

@@ -0,0 +1,166 @@
+"""
+LangGraph Agent that plans and executes a task step by step.
+
+Usage:
+    python agent.py "Compare Python and JavaScript"
+
+Dependencies:
+    - langgraph
+    - langchain-openai
+    - langchain-ollama (optional)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import sys
+import traceback
+from typing import TypedDict
+
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.graph import StateGraph, START, END
+from langchain_openai import ChatOpenAI
+
+# ----- State definition (from snippet) -----
+class PlanningState(TypedDict):
+    task: str
+    plan: list[str] | None
+    current_step: int
+    results: list[str]
+
+# ----- LLM configuration -----
+# Prefer OpenAI if API key is set; otherwise fallback to Ollama
+if os.getenv("OPENAI_API_KEY"):
+    llm = ChatOpenAI(
+        model="gpt-4o-mini",
+        temperature=0,
+    )
+else:
+    # Fall back to local Ollama model if available
+    llm = ChatOpenAI(
+        model=os.getenv("CHAT_MODEL", "llama3"),
+        base_url=os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1"),
+        api_key="ollama",
+        temperature=0,
+    )
+
+# ----- Planning node -----
+
+def planning(state: PlanningState) -> PlanningState:
+    """Prompt LLM to break the task into 3–6 numbered steps."""
+    try:
+        prompt = (
+            f"Given the task '{state['task']}', break it into 3-6 numbered steps. "
+            "Return only the numbered list or a JSON array of steps."
+        )
+        raw = llm.invoke(prompt)
+        text = raw if isinstance(raw, str) else raw.content
+    except Exception as e:
+        raise RuntimeError(f"LLM failed in planning node: {e}")
+
+    # Parse plan – first try JSON, then regex
+    plan: list[str] | None = None
+    try:
+        plan = json.loads(text)
+        if not isinstance(plan, list):
+            plan = None
+    except Exception:
+        plan = None
+
+    if plan is None:
+        pattern = r"^\s*\d+\.\s+(.*)$"
+        plan = [m.group(1).strip() for m in re.finditer(pattern, text, re.MULTILINE)]
+
+    if not plan or not (3 <= len(plan) <= 6):
+        raise ValueError(
+            f"Planning node returned invalid plan: {plan}. Expected 3-6 steps."
+        )
+
+    state["plan"] = plan
+    state["current_step"] = 0
+    state["results"] = []
+    return state
+
+# ----- Execution node -----
+
+def execution(state: PlanningState) -> PlanningState:
+    """Execute a single step from the plan and record the result."""
+    step_idx = state["current_step"]
+    if state["plan"] is None or step_idx >= len(state["plan"]):
+        return state
+    step = state["plan"][step_idx]
+    try:
+        prompt = f"Execute step: {step}\nProvide a concise result."
+        raw = llm.invoke(prompt)
+        result = raw if isinstance(raw, str) else raw.content
+    except Exception as e:
+        result = f"Error executing step: {e}"
+    state["results"].append(result.strip())
+    state["current_step"] = step_idx + 1
+    return state
+
+# ----- Graph construction -----
+# Graph diagram (from snippet):
+# START → planning → execution → should_continue
+#                      ↑____________| (execute)
+#                      finish → END
+
+builder = StateGraph(PlanningState)
+builder.add_node("planning", planning)
+builder.add_node("execution", execution)
+
+# Conditional edges to loop execution until all steps processed
+builder.add_conditional_edges(
+    "planning",
+    lambda state: "execution" if state.get("plan") else "END"
+)
+builder.add_conditional_edges(
+    "execution",
+    lambda state: "execution"
+    if state.get("current_step", 0) < len(state.get("plan", []))
+    else "END"
+)
+
+builder.set_entry_point("planning")
+
+# Compile graph with in-memory checkpointing
+graph = builder.compile(checkpointer=InMemorySaver())
+
+# ----- Demo -----
+
+def run_demo(task: str) -> None:
+    initial_state: PlanningState = {
+        "task": task,
+        "plan": None,
+        "current_step": 0,
+        "results": [],
+    }
+    try:
+        result = graph.invoke(initial_state)
+    except Exception:
+        traceback.print_exc()
+        sys.exit(1)
+
+    plan = result.get("plan", [])
+    print("\n===== PLAN =====")
+    for idx, step in enumerate(plan, 1):
+        print(f"{idx}. {step}")
+    print("\n===== RESULTS =====")
+    for idx, res in enumerate(result.get("results", []), 1):
+        print(f"[Step {idx}] {res}\n")
+    print("===== SUMMARY =====")
+    try:
+        summary_prompt = f"Given the collected results: {result.get('results', [])}, produce a concise summary of the task outcome."
+        summary = llm.invoke(summary_prompt)
+        print(summary if isinstance(summary, str) else summary.content)
+    except Exception as e:
+        print(f"Error generating summary: {e}")
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python agent.py '<task>'")
+        sys.exit(1)
+    task = " ".join(sys.argv[1:])
+    run_demo(task)