254 lines
8.4 KiB
Python
254 lines
8.4 KiB
Python
"""
|
|
CAD Assistant
|
|
|
|
Usage:
|
|
python main.py
|
|
"""
|
|
|
|
import asyncio
|
|
import sys
|
|
from time import sleep
|
|
|
|
from bootstrap_env import load_project_env
|
|
|
|
load_project_env()
|
|
|
|
from rich.console import Console
|
|
from rich.panel import Panel
|
|
|
|
from agent import get_agent, BaseAgent
|
|
from context.conversation_manager import get_conversation_manager
|
|
from observability import propagate_conversation_session
|
|
from react_stream import extract_output_text
|
|
|
|
console = Console()
|
|
|
|
|
|
def setup_agent() -> BaseAgent | None:
|
|
"""
|
|
Set up the agent.
|
|
"""
|
|
try:
|
|
# Use the global Agent singleton.
|
|
agent = get_agent(
|
|
model_name="cadagent",
|
|
)
|
|
console.print("CAD Assistant initialized successfully!")
|
|
return agent
|
|
except Exception as e:
|
|
print(f"Failed to initialize agent: {e}")
|
|
return None
|
|
|
|
|
|
def get_input() -> str:
|
|
"""
|
|
Get user input.
|
|
"""
|
|
lines = []
|
|
console.print("\n===========================")
|
|
console.print(">>> ", end="")
|
|
try:
|
|
while True:
|
|
line = input()
|
|
lines.append(line)
|
|
except EOFError:
|
|
pass
|
|
except KeyboardInterrupt:
|
|
return ""
|
|
return "\n".join(lines).strip()
|
|
|
|
|
|
def handle_special_commands(agent: BaseAgent, query: str) -> bool:
|
|
"""Handle special commands; return True if handled, otherwise False to continue normal processing."""
|
|
query_lower = query.lower().strip()
|
|
|
|
if query_lower == "/help":
|
|
console.print(
|
|
Panel.fit(
|
|
"[bold cyan]Special Commands:[/bold cyan]\n"
|
|
"[yellow]/help[/yellow] - Show this help\n"
|
|
"[yellow]/pad[/yellow] - Show SketchPad contents\n"
|
|
"[yellow]/pad_search <query>[/yellow] - Search SketchPad\n"
|
|
"[yellow]/pad_get <key>[/yellow] - Get content from SketchPad\n"
|
|
"[yellow]quit[/yellow] - Exit"
|
|
)
|
|
)
|
|
return True
|
|
|
|
if query_lower == "/pad":
|
|
try:
|
|
# Use BaseAgent._get_sketch_pad_summary to retrieve the summary.
|
|
summary = agent.get_sketch_pad_summary()
|
|
if summary.strip():
|
|
console.print(
|
|
Panel.fit(
|
|
summary,
|
|
title="[ SketchPad Contents ]",
|
|
border_style="cyan",
|
|
)
|
|
)
|
|
else:
|
|
console.print("[yellow]SketchPad is empty.[/yellow]")
|
|
except Exception as e:
|
|
console.print(f"[red]Unable to access SketchPad: {e}[/red]")
|
|
return True
|
|
|
|
if query_lower.startswith("/pad_search "):
|
|
search_query = query[12:].strip()
|
|
if search_query:
|
|
try:
|
|
results = agent.search_sketch_pad(search_query, 5)
|
|
if results:
|
|
formatted_lines = []
|
|
for item in results:
|
|
if isinstance(item, tuple) and len(item) >= 2:
|
|
key, sketch_item = item[0], item[1]
|
|
snippet = getattr(sketch_item, "summary", None) or str(
|
|
getattr(sketch_item, "value", "")
|
|
)
|
|
else:
|
|
key = item.get("key", "Unknown")
|
|
snippet = item.get("snippet", "No summary")
|
|
|
|
formatted_lines.append(f"- {key}: {str(snippet)[:50]}...")
|
|
|
|
content = "\n".join(formatted_lines)
|
|
console.print(
|
|
Panel.fit(
|
|
content,
|
|
title=f"[ SketchPad Search: '{search_query}' ]",
|
|
border_style="magenta",
|
|
)
|
|
)
|
|
else:
|
|
console.print(
|
|
f"[yellow]No SketchPad items found for '{search_query}'[/yellow]"
|
|
)
|
|
except Exception as e:
|
|
console.print(f"[red]Search failed: {e}[/red]")
|
|
else:
|
|
console.print("[red]Please provide a search query[/red]")
|
|
return True
|
|
|
|
if query_lower.startswith("/pad_get "):
|
|
key = query[9:].strip()
|
|
if key:
|
|
try:
|
|
value = agent.get_from_sketch_pad(key)
|
|
if value is not None:
|
|
# Truncate long content.
|
|
display_value = str(value)
|
|
if len(display_value) > 500:
|
|
display_value = display_value[:500] + "..."
|
|
|
|
console.print(
|
|
Panel.fit(
|
|
display_value,
|
|
title=f"[ SketchPad Item: {key} ]",
|
|
border_style="green",
|
|
)
|
|
)
|
|
else:
|
|
console.print(
|
|
f"[yellow]Key '{key}' not found in SketchPad[/yellow]"
|
|
)
|
|
except Exception as e:
|
|
console.print(f"[red]Get failed: {e}[/red]")
|
|
else:
|
|
console.print("[red]Usage: /pad_get <key>[/red]")
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
async def main() -> None:
|
|
"""
|
|
Main function.
|
|
"""
|
|
agent = setup_agent()
|
|
if not agent:
|
|
return
|
|
|
|
# Create a new conversation context.
|
|
conversation_manager = get_conversation_manager()
|
|
conversation = conversation_manager.create_conversation()
|
|
|
|
console.print(
|
|
Panel.fit(
|
|
f"[bold green]Ready![/bold green] Started new conversation session: [yellow]{conversation.uuid[:8]}...[/yellow]\n"
|
|
"[dim]Previous conversations are saved but not loaded automatically.[/dim]\n"
|
|
"[yellow]Create a new line and press [bold]Ctrl+D[/bold] (or [bold]Ctrl+Z[/bold] on Windows) to submit your query.[/yellow]\n"
|
|
"[cyan]Input 'quit' to exit the program.[/cyan]\n"
|
|
"[dim]Type '/help' for special commands, '/full_history' to view saved history.[/dim]",
|
|
title="[ CAD Assistant ]",
|
|
border_style="blue",
|
|
)
|
|
)
|
|
|
|
# Run within the conversation context.
|
|
with conversation:
|
|
while True:
|
|
try:
|
|
query = get_input()
|
|
if not query:
|
|
continue
|
|
|
|
if query.lower() == "quit":
|
|
break
|
|
|
|
# Handle special commands.
|
|
if handle_special_commands(agent, query):
|
|
continue
|
|
|
|
try:
|
|
console.print("===========================")
|
|
console.print("[Agent] >>> ", end="")
|
|
|
|
with propagate_conversation_session(
|
|
conversation_id=conversation.uuid,
|
|
metadata={
|
|
"model": getattr(agent, "model_name", None),
|
|
"agent_name": getattr(agent, "name", None),
|
|
"transport": "cli",
|
|
},
|
|
tags=["cadagent", "cli"],
|
|
):
|
|
async for output in agent.run(query):
|
|
try:
|
|
delta = extract_output_text(output, "cli")
|
|
except Exception:
|
|
delta = ""
|
|
if not delta:
|
|
continue
|
|
for char in delta:
|
|
if char == "\r":
|
|
char = "\n"
|
|
if char.strip() == "" and char not in ("\n", " "):
|
|
continue
|
|
console.print(char, end="")
|
|
sleep(0.01)
|
|
console.print("\n===========================")
|
|
except Exception as e:
|
|
console.print(f"\nError: {e}")
|
|
|
|
except KeyboardInterrupt:
|
|
console.print("\nType 'quit' to exit.")
|
|
continue
|
|
except Exception as e:
|
|
console.print(f"Error: {e}")
|
|
continue
|
|
|
|
# Save the conversation.
|
|
await conversation_manager.save_conversation(conversation.uuid)
|
|
console.print(f"[dim]Conversation {conversation.uuid[:8]}... saved.[/dim]")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
print("\nGoodbye!")
|
|
except Exception as e:
|
|
print(f"Fatal error: {e}")
|
|
sys.exit(1)
|