Skip to content

Commit cfa639f

Browse files
MarktechpostMarktechpost
authored
Add files via upload
1 parent f2db485 commit cfa639f

1 file changed

Lines changed: 281 additions & 0 deletions

File tree

‎AI Agents Codes/context_folding_llm_agent_long_horizon_Marktechpost.ipynb‎

Lines changed: 281 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,281 @@
1+
{
2+
"nbformat": 4,
3+
"nbformat_minor": 0,
4+
"metadata": {
5+
"colab": {
6+
"provenance": []
7+
},
8+
"kernelspec": {
9+
"name": "python3",
10+
"display_name": "Python 3"
11+
},
12+
"language_info": {
13+
"name": "python"
14+
}
15+
},
16+
"cells": [
17+
{
18+
"cell_type": "code",
19+
"source": [
20+
"import os, re, sys, math, random, json, textwrap, subprocess, shutil, time\n",
21+
"from typing import List, Dict, Tuple\n",
22+
"try:\n",
23+
" import transformers\n",
24+
"except:\n",
25+
" subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"transformers\", \"accelerate\", \"sentencepiece\"], check=True)\n",
26+
"from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline\n",
27+
"MODEL_NAME = os.environ.get(\"CF_MODEL\", \"google/flan-t5-small\")\n",
28+
"tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
29+
"model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\n",
30+
"llm = pipeline(\"text2text-generation\", model=model, tokenizer=tokenizer, device_map=\"auto\")\n",
31+
"def llm_gen(prompt: str, max_new_tokens=160, temperature=0.0) -> str:\n",
32+
" out = llm(prompt, max_new_tokens=max_new_tokens, do_sample=temperature>0.0, temperature=temperature)[0][\"generated_text\"]\n",
33+
" return out.strip()"
34+
],
35+
"metadata": {
36+
"colab": {
37+
"base_uri": "https://localhost:8080/"
38+
},
39+
"id": "5e7PBMlWvWh0",
40+
"outputId": "4b7b78bb-412b-4705-ca9e-da0fa75551bc"
41+
},
42+
"execution_count": 3,
43+
"outputs": [
44+
{
45+
"output_type": "stream",
46+
"name": "stderr",
47+
"text": [
48+
"Device set to use cpu\n"
49+
]
50+
}
51+
]
52+
},
53+
{
54+
"cell_type": "code",
55+
"source": [
56+
"import ast, operator as op\n",
57+
"OPS = {ast.Add: op.add, ast.Sub: op.sub, ast.Mult: op.mul, ast.Div: op.truediv, ast.Pow: op.pow, ast.USub: op.neg, ast.FloorDiv: op.floordiv, ast.Mod: op.mod}\n",
58+
"def _eval_node(n):\n",
59+
" if isinstance(n, ast.Num): return n.n\n",
60+
" if isinstance(n, ast.UnaryOp) and type(n.op) in OPS: return OPS[type(n.op)](_eval_node(n.operand))\n",
61+
" if isinstance(n, ast.BinOp) and type(n.op) in OPS: return OPS[type(n.op)](_eval_node(n.left), _eval_node(n.right))\n",
62+
" raise ValueError(\"Unsafe expression\")\n",
63+
"def calc(expr: str):\n",
64+
" node = ast.parse(expr, mode='eval').body\n",
65+
" return _eval_node(node)\n",
66+
"class FoldingMemory:\n",
67+
" def __init__(self, max_chars:int=800):\n",
68+
" self.active=[]; self.folds=[]; self.max_chars=max_chars\n",
69+
" def add(self,text:str):\n",
70+
" self.active.append(text.strip())\n",
71+
" while len(self.active_text())>self.max_chars and len(self.active)>1:\n",
72+
" popped=self.active.pop(0)\n",
73+
" fold=f\"- Folded: {popped[:120]}...\"\n",
74+
" self.folds.append(fold)\n",
75+
" def fold_in(self,summary:str): self.folds.append(summary.strip())\n",
76+
" def active_text(self)->str: return \"\\n\".join(self.active)\n",
77+
" def folded_text(self)->str: return \"\\n\".join(self.folds)\n",
78+
" def snapshot(self)->Dict: return {\"active_chars\":len(self.active_text()),\"n_folds\":len(self.folds)}"
79+
],
80+
"metadata": {
81+
"id": "YMrg5FjUwnrx"
82+
},
83+
"execution_count": 4,
84+
"outputs": []
85+
},
86+
{
87+
"cell_type": "code",
88+
"source": [
89+
"SUBTASK_DECOMP_PROMPT=\"\"\"You are an expert planner. Decompose the task below into 2-4 crisp subtasks.\n",
90+
"Return each subtask as a bullet starting with '- ' in priority order.\n",
91+
"Task: \"{task}\" \"\"\"\n",
92+
"SUBTASK_SOLVER_PROMPT=\"\"\"You are a precise problem solver with minimal steps.\n",
93+
"If a calculation is needed, write one line 'CALC(expr)'.\n",
94+
"Otherwise write 'ANSWER: <final>'.\n",
95+
"Think briefly; avoid chit-chat.\n",
96+
"\n",
97+
"Task: {task}\n",
98+
"Subtask: {subtask}\n",
99+
"Notes (folded context):\n",
100+
"{notes}\n",
101+
"\n",
102+
"Now respond with either CALC(...) or ANSWER: ...\"\"\"\n",
103+
"SUBTASK_SUMMARY_PROMPT=\"\"\"Summarize the subtask outcome in <=3 bullets, total <=50 tokens.\n",
104+
"Subtask: {name}\n",
105+
"Steps:\n",
106+
"{trace}\n",
107+
"Final: {final}\n",
108+
"Return only bullets starting with '- '.\"\"\"\n",
109+
"FINAL_SYNTH_PROMPT=\"\"\"You are a senior agent. Synthesize a final, coherent solution using ONLY:\n",
110+
"- The original task\n",
111+
"- Folded summaries (below)\n",
112+
"Avoid repeating steps. Be concise and actionable.\n",
113+
"\n",
114+
"Task: {task}\n",
115+
"Folded summaries:\n",
116+
"{folds}\n",
117+
"\n",
118+
"Final answer:\"\"\"\n",
119+
"def parse_bullets(text:str)->List[str]:\n",
120+
" return [ln[2:].strip() for ln in text.splitlines() if ln.strip().startswith(\"- \")]"
121+
],
122+
"metadata": {
123+
"id": "082OEce2wqSL"
124+
},
125+
"execution_count": 5,
126+
"outputs": []
127+
},
128+
{
129+
"cell_type": "code",
130+
"source": [
131+
"def run_subtask(task:str, subtask:str, memory:FoldingMemory, max_tool_iters:int=3)->Tuple[str,str,List[str]]:\n",
132+
" notes=(memory.folded_text() or \"(none)\")\n",
133+
" trace=[]; final=\"\"\n",
134+
" for _ in range(max_tool_iters):\n",
135+
" prompt=SUBTASK_SOLVER_PROMPT.format(task=task,subtask=subtask,notes=notes)\n",
136+
" out=llm_gen(prompt,max_new_tokens=96); trace.append(out)\n",
137+
" m=re.search(r\"CALC\\((.+?)\\)\",out)\n",
138+
" if m:\n",
139+
" try:\n",
140+
" val=calc(m.group(1))\n",
141+
" trace.append(f\"TOOL:CALC -> {val}\")\n",
142+
" out2=llm_gen(prompt+f\"\\nTool result: {val}\\nNow produce 'ANSWER: ...' only.\",max_new_tokens=64)\n",
143+
" trace.append(out2)\n",
144+
" if out2.strip().startswith(\"ANSWER:\"):\n",
145+
" final=out2.split(\"ANSWER:\",1)[1].strip(); break\n",
146+
" except Exception as e:\n",
147+
" trace.append(f\"TOOL:CALC ERROR -> {e}\")\n",
148+
" if out.strip().startswith(\"ANSWER:\"):\n",
149+
" final=out.split(\"ANSWER:\",1)[1].strip(); break\n",
150+
" if not final:\n",
151+
" final=\"No definitive answer; partial reasoning:\\n\"+\"\\n\".join(trace[-2:])\n",
152+
" summ=llm_gen(SUBTASK_SUMMARY_PROMPT.format(name=subtask,trace=\"\\n\".join(trace),final=final),max_new_tokens=80)\n",
153+
" summary_bullets=\"\\n\".join(parse_bullets(summ)[:3]) or f\"- {subtask}: {final[:60]}...\"\n",
154+
" return final, summary_bullets, trace\n",
155+
"class ContextFoldingAgent:\n",
156+
" def __init__(self,max_active_chars:int=800):\n",
157+
" self.memory=FoldingMemory(max_chars=max_active_chars)\n",
158+
" self.metrics={\"subtasks\":0,\"tool_calls\":0,\"chars_saved_est\":0}\n",
159+
" def decompose(self,task:str)->List[str]:\n",
160+
" plan=llm_gen(SUBTASK_DECOMP_PROMPT.format(task=task),max_new_tokens=96)\n",
161+
" subs=parse_bullets(plan)\n",
162+
" return subs[:4] if subs else [\"Main solution\"]\n",
163+
" def run(self,task:str)->Dict:\n",
164+
" t0=time.time()\n",
165+
" self.memory.add(f\"TASK: {task}\")\n",
166+
" subtasks=self.decompose(task)\n",
167+
" self.metrics[\"subtasks\"]=len(subtasks)\n",
168+
" folded=[]\n",
169+
" for st in subtasks:\n",
170+
" self.memory.add(f\"SUBTASK: {st}\")\n",
171+
" final,fold_summary,trace=run_subtask(task,st,self.memory)\n",
172+
" self.memory.fold_in(fold_summary)\n",
173+
" folded.append(f\"- {st}: {final}\")\n",
174+
" self.memory.add(f\"SUBTASK_DONE: {st}\")\n",
175+
" final=llm_gen(FINAL_SYNTH_PROMPT.format(task=task,folds=self.memory.folded_text()),max_new_tokens=200)\n",
176+
" t1=time.time()\n",
177+
" return {\"task\":task,\"final\":final.strip(),\"folded_summaries\":self.memory.folded_text(),\n",
178+
" \"active_context_chars\":len(self.memory.active_text()),\n",
179+
" \"subtask_finals\":folded,\"runtime_sec\":round(t1-t0,2)}"
180+
],
181+
"metadata": {
182+
"id": "pcTrP-5zwsXc"
183+
},
184+
"execution_count": 6,
185+
"outputs": []
186+
},
187+
{
188+
"cell_type": "code",
189+
"source": [
190+
"DEMO_TASKS=[\n",
191+
" \"Plan a 3-day study schedule for ML with daily workouts and simple meals; include time blocks.\",\n",
192+
" \"Compute a small project budget with 3 items (laptop 799.99, course 149.5, snacks 23.75), add 8% tax and 5% buffer, and present a one-paragraph recommendation.\"\n",
193+
"]\n",
194+
"def pretty(d): return json.dumps(d, indent=2, ensure_ascii=False)\n",
195+
"if __name__==\"__main__\":\n",
196+
" agent=ContextFoldingAgent(max_active_chars=700)\n",
197+
" for i,task in enumerate(DEMO_TASKS,1):\n",
198+
" print(\"=\"*70)\n",
199+
" print(f\"DEMO #{i}: {task}\")\n",
200+
" res=agent.run(task)\n",
201+
" print(\"\\n--- Folded Summaries ---\\n\"+(res[\"folded_summaries\"] or \"(none)\"))\n",
202+
" print(\"\\n--- Final Answer ---\\n\"+res[\"final\"])\n",
203+
" print(\"\\n--- Diagnostics ---\")\n",
204+
" diag={k:res[k] for k in [\"active_context_chars\",\"runtime_sec\"]}\n",
205+
" diag[\"n_subtasks\"]=len(agent.decompose(task))\n",
206+
" print(pretty(diag))"
207+
],
208+
"metadata": {
209+
"colab": {
210+
"base_uri": "https://localhost:8080/"
211+
},
212+
"id": "ty_CUdKTwuZh",
213+
"outputId": "5aa7008d-3bbc-462e-bcd5-87208f445bd5"
214+
},
215+
"execution_count": 7,
216+
"outputs": [
217+
{
218+
"output_type": "stream",
219+
"name": "stdout",
220+
"text": [
221+
"======================================================================\n",
222+
"DEMO #1: Plan a 3-day study schedule for ML with daily workouts and simple meals; include time blocks.\n"
223+
]
224+
},
225+
{
226+
"output_type": "stream",
227+
"name": "stderr",
228+
"text": [
229+
"/tmp/ipython-input-194531414.py:4: DeprecationWarning: ast.Num is deprecated and will be removed in Python 3.14; use ast.Constant instead\n",
230+
" if isinstance(n, ast.Num): return n.n\n"
231+
]
232+
},
233+
{
234+
"output_type": "stream",
235+
"name": "stdout",
236+
"text": [
237+
"\n",
238+
"--- Folded Summaries ---\n",
239+
"- Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule: No definitive answer; partial reasoning:\n",
240+
"CALC(...)\n",
241+
"TOOL:CALC...\n",
242+
"- Folded: TASK: Plan a 3-day study schedule for ML with daily workouts and simple meals; include time blocks....\n",
243+
"- Folded: SUBTASK: Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study sched...\n",
244+
"\n",
245+
"--- Final Answer ---\n",
246+
"Folded summaries (below) Avoid repeating steps. Be concise and actionable. Task: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for\n",
247+
"\n",
248+
"--- Diagnostics ---\n",
249+
"{\n",
250+
" \"active_context_chars\": 368,\n",
251+
" \"runtime_sec\": 24.1,\n",
252+
" \"n_subtasks\": 1\n",
253+
"}\n",
254+
"======================================================================\n",
255+
"DEMO #2: Compute a small project budget with 3 items (laptop 799.99, course 149.5, snacks 23.75), add 8% tax and 5% buffer, and present a one-paragraph recommendation.\n",
256+
"\n",
257+
"--- Folded Summaries ---\n",
258+
"- Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study schedule: No definitive answer; partial reasoning:\n",
259+
"CALC(...)\n",
260+
"TOOL:CALC...\n",
261+
"- Folded: TASK: Plan a 3-day study schedule for ML with daily workouts and simple meals; include time blocks....\n",
262+
"- Folded: SUBTASK: Tasks: Plan a 3-day study schedule for ML with daily workouts and simple meals. Tasks: Plan a 3-day study sched...\n",
263+
"- Main solution: No definitive answer; partial reasoning:\n",
264+
"CALC(...)\n",
265+
"TOOL:CALC...\n",
266+
"\n",
267+
"--- Final Answer ---\n",
268+
"No definitive answer; partial reasoning: CALC(...) TOOL:CALC...\n",
269+
"\n",
270+
"--- Diagnostics ---\n",
271+
"{\n",
272+
" \"active_context_chars\": 584,\n",
273+
" \"runtime_sec\": 10.34,\n",
274+
" \"n_subtasks\": 1\n",
275+
"}\n"
276+
]
277+
}
278+
]
279+
}
280+
]
281+
}

0 commit comments

Comments
 (0)