Spaces:
Sleeping
Sleeping
File size: 8,753 Bytes
f97126e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | import asyncio
import os
from collections import deque
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from rich.table import Table
from rich.text import Text
from rich.align import Align
from langchain_core.messages import HumanMessage, AIMessage
from src import retriever
from src.retriever import graph
from src.react_agent.agent import COMPILED_AGENT
from src.react_agent import generate
console = Console()
async def main():
console.print(Panel(
Align.center("[bold magenta]Vectorless-RAG ReAct Agent Assistant[/bold magenta]\n"
"[dim white]Autonomous Thought -> Action -> Observation loop[/dim white]"),
border_style="magenta"
))
console.print("[bold yellow]Loading indices...[/bold yellow]")
retriever.load("tree")
console.print()
console.print(Panel(
"[bold green]Ready![/bold green] Ask a legal scenario. The agent will autonomously decide what to search.\n\n"
"[bold white]Commands:[/bold white]\n"
" [cyan]exit[/cyan] / [cyan]quit[/cyan] - Close assistant\n"
" [cyan]clear[/cyan] - Clear chat memory\n"
" [cyan]trace[/cyan] - Toggle verbose reasoning trace (currently ON)",
title="[bold magenta]ReAct System Status[/bold magenta]",
border_style="magenta",
expand=False
))
console.print()
MEMORY_LIMIT = 5
history = deque(maxlen=MEMORY_LIMIT)
trace_mode = True
while True:
try:
query = console.input("[bold deep_sky_blue1]Query > [/bold deep_sky_blue1]")
if query.strip().lower() in ['exit', 'quit', 'q']:
break
if query.strip().lower() == 'clear':
history.clear()
console.print("[bold green]Memory cleared![/bold green]\n")
continue
if query.strip().lower() == 'trace':
trace_mode = not trace_mode
status = "ENABLED" if trace_mode else "DISABLED"
console.print(f"[bold yellow]Trace logs {status}.[/bold yellow]\n")
continue
if not query.strip():
continue
# If trace mode is ON, we stream updates to show the Thought-Action-Observation loop
if trace_mode:
console.print(f"\n[bold yellow]Agent is reasoning...[/bold yellow]")
# 1. Format history messages
messages = []
for turn in history:
messages.append(HumanMessage(content=turn.get("user", "")))
assistant_clean = turn.get("assistant", "").split("[References]")[0].strip()
messages.append(AIMessage(content=assistant_clean))
messages.append(HumanMessage(content=query))
# 2. Run streaming graph
try:
async for event in COMPILED_AGENT.astream(
{"messages": messages},
config={"recursion_limit": 10},
stream_mode="updates"
):
for node, update in event.items():
if node == "agent":
msgs = update.get("messages", [])
if msgs:
msg = msgs[-1]
# Handle list content in agent thoughts
content = msg.content
if isinstance(content, list):
parts = []
for part in content:
if isinstance(part, str):
parts.append(part)
elif isinstance(part, dict) and "text" in part:
parts.append(part["text"])
elif hasattr(part, "text"):
parts.append(part.text)
content = "".join(parts)
if content:
console.print(Panel(
content.strip(),
title="[bold yellow]Agent Thought[/bold yellow]",
border_style="yellow"
))
if hasattr(msg, "tool_calls") and msg.tool_calls:
for tc in msg.tool_calls:
console.print(f"[bold cyan]Action (Call Tool):[/bold cyan] [bold white]{tc['name']}[/bold white] with args: [magenta]{tc['args']}[/magenta]")
elif node == "tools":
msgs = update.get("messages", [])
if msgs:
msg = msgs[-1]
# Handle list content in tool observations
content = msg.content
if isinstance(content, list):
parts = []
for part in content:
if isinstance(part, str):
parts.append(part)
elif isinstance(part, dict) and "text" in part:
parts.append(part["text"])
elif hasattr(part, "text"):
parts.append(part.text)
content = "".join(parts)
preview = content[:300] + "..." if len(content) > 300 else content
console.print(Panel(
preview.strip(),
title="[bold green]Observation (Tool Output)[/bold green]",
border_style="green"
))
console.print()
except Exception as e:
console.print(f"\n[bold red]Trace Loop Error: {e}[/bold red]\n")
# 3. Call standard generate interface to get final formatted answer & metadata
with console.status("[bold yellow]Synthesizing final structured response...[/bold yellow]", spinner="dots"):
res = await generate(
query=query,
history=list(history),
)
# 4. Print Response Metadata (ASCII only to prevent Windows console encoding crash)
conf_badge = "[bold green][OK] ADEQUATE CONTEXT[/bold green]" if res["confidence"] > 0 else "[bold red][FAIL] INSUFFICIENT CONTEXT[/bold red]"
border_color = "green" if res["confidence"] > 0 else "red"
console.print(Panel(
f"Status: {conf_badge}\nLatency: [cyan]{res['latency_ms']}[/cyan] ms",
title="[bold white]Response Metadata[/bold white]",
border_style=border_color,
expand=False
))
console.print()
# 5. Print Answer using Markdown
console.print(Panel(
Markdown(res.get('answer', '')),
title="[bold magenta]ReAct Final Answer[/bold magenta]",
border_style="magenta"
))
console.print()
# Add to memory
history.append({
"user": query,
"assistant": res.get("answer", "")
})
console.print("[dim white]" + "="*60 + "[/dim white]\n")
except KeyboardInterrupt:
break
except Exception as e:
console.print(f"\n[bold red]Error: {e}[/bold red]\n")
if __name__ == "__main__":
asyncio.run(main())
|