diff --git a/CLI.md b/CLI.md new file mode 100644 index 0000000..6b982db --- /dev/null +++ b/CLI.md @@ -0,0 +1,140 @@ +# CLI Usage Guide + +## Interactive CLI (`cli.py`) + +Antarmuka terminal interaktif untuk Perplexity AI dengan multi-turn history, streaming, dan MCP tool integration. + +### Setup + +1. Clone repo dan install dependencies: + ```bash + git clone https://github.com/ardzz/perplexity-scrape.git + cd perplexity-scrape + pip install -r requirements.txt + ``` + +2. Copy `.env.example` dan isi credential Perplexity: + ```bash + cp .env.example .env + # Edit .env dan isi PERPLEXITY_SESSION_TOKEN, dll + ``` + +### Menjalankan CLI + +```bash +# Mode default (claude46sonnetthinking) +python3 cli.py + +# Pilih model spesifik +python3 cli.py --model sonar-pro +python3 cli.py --model claude46sonnetthinking + +# Dengan MCP tools (BurpSuite, dll) +python3 cli.py --mcp ./mcp_tools.json + +# Load sesi sebelumnya +python3 cli.py --load ~/.perplexity_cli/sessions/session_20240525.json + +# Mode incognito (tidak muncul di dashboard Perplexity) +python3 cli.py --incognito + +# System prompt kustom +python3 cli.py --system "Kamu adalah security researcher. Fokus ke vulnerability analysis." + +# Kombinasi lengkap +python3 cli.py --model claude46sonnetthinking --mcp ./mcp_tools.json --incognito +``` + +### Slash Commands + +| Command | Fungsi | +|---|---| +| `/model ` | Ganti model aktif | +| `/models` | List semua model yang tersedia | +| `/clear` | Hapus history percakapan sesi ini | +| `/save [path]` | Simpan sesi ke file JSON | +| `/load ` | Load sesi dari file JSON | +| `/sessions` | List sesi tersimpan | +| `/system ` | Set system prompt | +| `/system` | Tampilkan system prompt aktif | +| `/mode ` | Ganti mode: `copilot`, `search`, `reasoning`, `writing` | +| `/focus ` | Ganti fokus: `internet`, `academic`, `wolfram`, `youtube`, `reddit` | +| `/incognito` | Toggle incognito mode | +| `/tools` | List semua MCP tools yang terhubung | +| `/tool [args]` | Panggil MCP tool secara manual | +| `/mcp ` | Connect ke MCP servers dari config file | +| `/history` | Tampilkan riwayat percakapan | +| `/status` | Info sesi aktif | +| `/help` | Bantuan | +| `exit` / `quit` | Keluar | + +### MCP Tool Integration + +Buat file `mcp_tools.json` (lihat `mcp_tools.example.json`): + +```json +{ + "mcpServers": { + "burpsuite": { + "command": "python", + "args": ["/path/to/burpsuite-mcp/server.py"], + "env": { + "BURP_API_KEY": "your-key" + } + }, + "filesystem": { + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/home/user"] + } + } +} +``` + +Jalankan dengan MCP: +```bash +python3 cli.py --mcp ./mcp_tools.json +``` + +Saat ada tools terhubung, LLM bisa call tools dengan format: +``` +TOOL_CALL: burpsuite.active_scan {"url": "https://target.com", "type": "xss"} +``` +CLI akan otomatis mengeksekusi dan mengembalikan hasil ke konteks percakapan. + +Atau panggil manual: +``` +/tool burpsuite.active_scan url=https://target.com type=xss +``` + +### Model yang Tersedia + +- `claude46sonnetthinking` (default) +- `claude45sonnetthinking` +- `claude37sonnetthinking` +- `claude35sonnet` +- `sonar-pro` +- `sonar` +- `sonar-reasoning-pro` +- `sonar-reasoning` +- `r1-1776` +- `gpt-4o` +- `o3-mini` +- `gemini-2.0-flash` + +### Session Management + +Sesi disimpan otomatis di `~/.perplexity_cli/sessions/` saat keluar. + +List sesi: +```bash +# Via CLI +/sessions + +# Atau langsung +ls ~/.perplexity_cli/sessions/ +``` + +Load sesi lama: +```bash +python3 cli.py --load ~/.perplexity_cli/sessions/session_20240525_143022.json +``` diff --git a/cli.py b/cli.py new file mode 100644 index 0000000..88a6201 --- /dev/null +++ b/cli.py @@ -0,0 +1,810 @@ +#!/usr/bin/env python3 +""" +Perplexity Interactive CLI + +An interactive terminal interface for Perplexity AI with: +- Multi-turn conversation history +- Streaming real-time output +- MCP tool integration (BurpSuite, filesystem, custom servers) +- Rich markdown rendering + +Usage: + python3 cli.py + python3 cli.py --model sonar-pro + python3 cli.py --mcp ./mcp_tools.json + python3 cli.py --load ~/sessions/chat.json +""" + +import argparse +import json +import os +import sys +import time +import signal +from pathlib import Path +from typing import Optional + +# Ensure project root is in path +sys.path.insert(0, str(Path(__file__).parent)) + +from dotenv import load_dotenv +load_dotenv() + +from rich.console import Console +from rich.live import Live +from rich.spinner import Spinner +from rich.text import Text +from prompt_toolkit import PromptSession +from prompt_toolkit.history import FileHistory +from prompt_toolkit.auto_suggest import AutoSuggestFromHistory +from prompt_toolkit.styles import Style +from prompt_toolkit.formatted_text import HTML + +from src.core.perplexity_client import PerplexityClient, PerplexityResponse +from src.cli.session import ChatSession, AVAILABLE_MODELS, DEFAULT_MODEL, _ALL_MODELS, get_model_name +from src.cli.renderer import Renderer, console +from src.cli.mcp_bridge import MCPToolBridge + + +# ─── Prompt style ──────────────────────────────────────────────────────────── + +PROMPT_STYLE = Style.from_dict({ + "prompt": "#00d7ff bold", + "prompt.arrow": "#888888", +}) + + +# ─── Multi-strategy text extraction ───────────────────────────────────────── + +def _deep_find_text(obj, depth: int = 0, max_depth: int = 6) -> str: + """Recursively search a dict/list for the longest meaningful text string.""" + if depth > max_depth: + return "" + best = "" + + if isinstance(obj, str): + cleaned = obj.strip() + if len(cleaned) > len(best) and len(cleaned) > 20: + best = cleaned + # Try to parse as JSON too + if cleaned.startswith(("{", "[")): + try: + parsed = json.loads(cleaned) + candidate = _deep_find_text(parsed, depth + 1, max_depth) + if len(candidate) > len(best): + best = candidate + except (json.JSONDecodeError, TypeError): + pass + elif isinstance(obj, dict): + for k, v in obj.items(): + if k in ("url", "query", "id", "uuid", "sha", "hash", "frontend_uuid", + "context_uuid", "request_id", "visitor_id"): + continue + candidate = _deep_find_text(v, depth + 1, max_depth) + if len(candidate) > len(best): + best = candidate + elif isinstance(obj, list): + for item in obj: + candidate = _deep_find_text(item, depth + 1, max_depth) + if len(candidate) > len(best): + best = candidate + + return best + + +def _extract_text_from_events(events: list[dict]) -> tuple[str, list[dict]]: + """ + Multi-strategy text extraction from Perplexity SSE events. + Handles different response formats across models (claude, gpt-4o, sonar, r1, etc.) + + Returns (text, citations). + """ + text = "" + citations: list[dict] = [] + + # ── Strategy 1: streaming diff blocks (claude/sonar models) ─────────── + buffer = "" + for event in events: + for block in event.get("blocks", []): + usage = block.get("intended_usage", "") + if usage in ("ask_text_0_markdown", "ask_text", "answer"): + for patch in block.get("diff_block", {}).get("patches", []): + op = patch.get("op") + value = patch.get("value") + if op == "replace" and isinstance(value, dict): + buffer = "".join(value.get("chunks", [])) + elif op == "add": + if isinstance(value, str): + buffer += value + elif isinstance(value, dict): + buffer += "".join(value.get("chunks", [])) + if buffer.strip(): + text = buffer.strip() + + # ── Strategy 2: FINAL event nested JSON (all models) ────────────────── + for event in events: + if event.get("step_type") != "FINAL": + continue + + text_field = event.get("text", "") + if not text_field: + # Some models put the answer directly in other fields + for field in ("answer", "output", "response", "content"): + val = event.get(field, "") + if isinstance(val, str) and val.strip(): + if not text: + text = val.strip() + continue + + # text_field might be a JSON array of step objects OR plain text + if isinstance(text_field, str) and text_field.strip().startswith("["): + try: + steps = json.loads(text_field) + for step in steps: + step_content = step.get("content", {}) + inner_type = step.get("step_type", "") + + # Citations from SEARCH_RESULTS step + if inner_type == "SEARCH_RESULTS": + for wr in step_content.get("web_results", []): + if wr.get("name") and wr.get("url"): + citations.append({ + "title": wr.get("name", ""), + "url": wr.get("url", ""), + "snippet": wr.get("snippet", ""), + }) + + # Answer from inner FINAL step + if inner_type == "FINAL": + answer_str = step_content.get("answer", "") + if answer_str: + try: + answer_data = json.loads(answer_str) + if "answer" in answer_data and not text: + text = answer_data["answer"] + # structured_answer overrides plain answer + for item in answer_data.get("structured_answer", []): + if item.get("type") == "markdown" and item.get("text"): + text = item["text"] + # Additional citations from answer's web_results + for wr in answer_data.get("web_results", []): + if wr.get("name") and wr.get("url"): + citations.append({ + "title": wr.get("name", ""), + "url": wr.get("url", ""), + "snippet": wr.get("snippet", ""), + }) + except (json.JSONDecodeError, AttributeError): + # answer_str is plain text already + if not text and answer_str.strip(): + text = answer_str.strip() + + # structured_answer at step level + if not text: + for item in step_content.get("structured_answer", []): + if item.get("type") == "markdown" and item.get("text"): + text = item["text"] + break + + except (json.JSONDecodeError, TypeError, ValueError): + # Not valid JSON — treat text_field as plain text + if not text and isinstance(text_field, str) and text_field.strip(): + text = text_field.strip() + + elif isinstance(text_field, str) and text_field.strip(): + # Plain text in text_field (some model formats) + if not text: + text = text_field.strip() + + # ── Strategy 3: direct top-level fields in any event ────────────────── + if not text: + for event in events: + for field in ("answer", "output", "response", "content", "message", "text"): + val = event.get(field, "") + if isinstance(val, str) and val.strip() and len(val.strip()) > 10: + # Ignore if it looks like JSON (handled above) + stripped = val.strip() + if not stripped.startswith(("[", "{")): + text = stripped + break + if text: + break + + # ── Strategy 4: recursive deep scan of FINAL event ──────────────────── + if not text: + for event in events: + if event.get("step_type") == "FINAL": + candidate = _deep_find_text(event) + if len(candidate) > 20: + text = candidate + break + + # Deduplicate citations + seen_urls = set() + unique_citations = [] + for c in citations: + if c["url"] not in seen_urls: + seen_urls.add(c["url"]) + unique_citations.append(c) + + return text, unique_citations + + +# ─── Streaming handler ─────────────────────────────────────────────────────── + +def stream_perplexity( + client: PerplexityClient, + query: str, + model: str, + mode: str, + search_focus: str, + is_incognito: bool, + renderer: Renderer, +) -> PerplexityResponse: + """ + Stream a query to Perplexity and render output in real time. + Uses multi-strategy text extraction to support all models. + Returns the complete PerplexityResponse when done. + """ + result = PerplexityResponse() + text_buffer = "" + start_time = time.time() + + with renderer.start_assistant_stream(model) as live: + try: + for event in client.ask_stream( + query=query, + mode=mode, + model_preference=model, + search_focus=search_focus, + is_incognito=is_incognito, + ): + result.raw_events.append(event) + + # Live preview: track streaming chunks for spinner feedback + for block in event.get("blocks", []): + usage = block.get("intended_usage", "") + if usage in ("ask_text_0_markdown", "ask_text", "answer"): + for patch in block.get("diff_block", {}).get("patches", []): + op = patch.get("op") + value = patch.get("value") + if op == "replace" and isinstance(value, dict): + text_buffer = "".join(value.get("chunks", [])) + elif op == "add": + if isinstance(value, str): + text_buffer += value + elif isinstance(value, dict): + text_buffer += "".join(value.get("chunks", [])) + + if text_buffer: + spinner_text = Text() + spinner_text.append(" Streaming ", style="dim") + spinner_text.append(f"({len(text_buffer)} chars)", style="dim bright_cyan") + live.update(Spinner("dots2", text=spinner_text, style="bright_cyan")) + + except KeyboardInterrupt: + raise + except Exception as e: + renderer.print_error(f"Stream error: {e}") + + # ── Multi-strategy extraction ───────────────────────────────────────── + extracted_text, citations = _extract_text_from_events(result.raw_events) + + # Final fallback: use the live-streamed buffer + if not extracted_text and text_buffer.strip(): + extracted_text = text_buffer.strip() + + result.text = extracted_text + result.citations = citations + + elapsed = time.time() - start_time + + if result.text: + renderer.print_assistant_response( + text=result.text, + citations=result.citations, + model=model, + elapsed=elapsed, + ) + else: + renderer.print_assistant_response( + text="*(no response — model may not be available or session token expired)*", + citations=[], + model=model, + elapsed=elapsed, + ) + renderer.print_warning( + f"Got {len(result.raw_events)} SSE events but could not extract text. " + "Check your session token or try a different model." + ) + + return result + + +# ─── Command handlers ───────────────────────────────────────────────────────── + +def handle_command( + cmd: str, + session: ChatSession, + renderer: Renderer, + bridge: MCPToolBridge, +) -> Optional[str]: + """ + Handle a slash command. + Returns None if handled, or 'exit' to quit. + """ + parts = cmd.strip().split(None, 1) + command = parts[0].lower() + arg = parts[1].strip() if len(parts) > 1 else "" + + if command in ("/help", "/?"): + renderer.print_help() + + elif command == "/model": + if not arg: + renderer.print_models_table(_ALL_MODELS, session.model) + else: + old = session.model + session.set_model(arg) + display = get_model_name(arg) + renderer.print_success(f"Model changed: [dim]{old}[/] → [bold bright_cyan]{arg}[/] [dim]({display})[/]") + + elif command == "/models": + renderer.print_models_table(_ALL_MODELS, session.model) + + elif command == "/clear": + session.clear() + renderer.print_cleared() + + elif command == "/save": + path = session.save(arg if arg else None) + renderer.print_success(f"Session saved to: [dim]{path}[/]") + + elif command == "/load": + if not arg: + renderer.print_error("Usage: /load ") + else: + try: + loaded = ChatSession.load(arg) + session.__dict__.update(loaded.__dict__) + renderer.print_success( + f"Session loaded: [dim]{loaded.session_id}[/] " + f"({loaded.message_count} messages)" + ) + except Exception as e: + renderer.print_error(f"Failed to load: {e}") + + elif command == "/sessions": + sessions = ChatSession.list_saved_sessions() + renderer.print_sessions_table(sessions) + + elif command == "/system": + if arg: + session.set_system_prompt(arg) + renderer.print_success("System prompt updated.") + else: + if session.system_prompt: + renderer.print_info(f"System prompt: [dim]{session.system_prompt}[/]") + else: + renderer.print_info("No system prompt set.") + + elif command == "/mode": + VALID_MODES = ["copilot", "search", "reasoning", "writing"] + if arg in VALID_MODES: + session.mode = arg + renderer.print_success(f"Mode set to: [bold bright_cyan]{arg}[/]") + else: + renderer.print_error(f"Invalid mode. Choose from: {', '.join(VALID_MODES)}") + + elif command == "/focus": + VALID_FOCUS = ["internet", "academic", "writing", "wolfram", "youtube", "reddit"] + if arg in VALID_FOCUS: + session.search_focus = arg + renderer.print_success(f"Search focus set to: [bold bright_cyan]{arg}[/]") + else: + renderer.print_error(f"Invalid focus. Choose from: {', '.join(VALID_FOCUS)}") + + elif command == "/incognito": + session.is_incognito = not session.is_incognito + state = "ON 🕵" if session.is_incognito else "OFF" + renderer.print_success(f"Incognito mode: [bold]{state}[/]") + + elif command == "/tools": + tools = bridge.get_all_tools() + renderer.print_tools_table(tools) + + elif command == "/tool": + if not arg: + renderer.print_error("Usage: /tool [key=value ...]") + else: + tool_parts = arg.split(None, 1) + tool_name = tool_parts[0] + tool_args_str = tool_parts[1] if len(tool_parts) > 1 else "" + try: + args = bridge.parse_tool_args(tool_args_str) + renderer.print_tool_call(tool_name, args) + result, elapsed = bridge.call_tool(tool_name, args) + renderer.print_tool_result(tool_name, result, elapsed) + session.add_tool_result(tool_name, result) + except KeyError as e: + renderer.print_error(f"Tool not found: {e}") + except Exception as e: + renderer.print_tool_error(tool_name, str(e)) + + elif command == "/mcp": + if not arg: + renderer.print_error("Usage: /mcp ") + else: + try: + loaded = bridge.load_config(arg) + renderer.print_info(f"Loaded {len(loaded)} server(s): {', '.join(loaded)}") + renderer.print_info("Connecting…") + results = bridge.connect_all() + for name, ok in results.items(): + if ok: + tools = [t for t in bridge.get_all_tools() if t["server"] == name] + renderer.print_success(f"[{name}] Connected ({len(tools)} tools)") + else: + renderer.print_error(f"[{name}] Connection failed") + except FileNotFoundError as e: + renderer.print_error(str(e)) + except Exception as e: + renderer.print_error(f"MCP load error: {e}") + + elif command == "/history": + renderer.print_history(session.get_history_for_display()) + + elif command == "/status": + renderer.print_status(session) + + elif command == "/debug": + # Show last N raw events for debugging + events = session.messages + n = int(arg) if arg.isdigit() else 3 + renderer.print_info(f"Last {n} raw SSE events from previous response:") + # This is a dev tool — users shouldn't normally need it + console.print("[dim]Use /debug to show raw event structures[/]") + + elif command in ("exit", "quit", "/exit", "/quit"): + return "exit" + + else: + renderer.print_error(f"Unknown command: [bold]{command}[/] — type [bold]/help[/] for help") + + return None + + +# ─── Main REPL ──────────────────────────────────────────────────────────────── + +def run_cli( + model: str = DEFAULT_MODEL, + mcp_config: Optional[str] = None, + load_session: Optional[str] = None, + mode: str = "copilot", + search_focus: str = "internet", + system_prompt: Optional[str] = None, + incognito: bool = False, +) -> None: + """Main CLI entry point.""" + + # ── Init client ─────────────────────────────────────── + try: + client = PerplexityClient() + except ValueError as e: + console.print(f"[bold red]Error:[/] {e}") + console.print( + "[dim]Create a .env file with PERPLEXITY_SESSION_TOKEN and other required vars.[/]" + ) + sys.exit(1) + + renderer = Renderer() + bridge = MCPToolBridge() + + # ── Load or create session ──────────────────────────── + if load_session: + try: + session = ChatSession.load(load_session) + renderer.print_success(f"Session loaded: {session.session_id}") + except Exception as e: + renderer.print_error(f"Could not load session: {e}") + session = ChatSession(model=model, mode=mode, search_focus=search_focus, + system_prompt=system_prompt, is_incognito=incognito) + else: + session = ChatSession( + model=model, + mode=mode, + search_focus=search_focus, + system_prompt=system_prompt, + is_incognito=incognito, + ) + + # ── Load MCP config ─────────────────────────────────── + connected_tools = [] + if mcp_config: + try: + loaded = bridge.load_config(mcp_config) + results = bridge.connect_all() + for name, ok in results.items(): + if ok: + tools = [t["full_name"] for t in bridge.get_all_tools() if t["server"] == name] + connected_tools.extend(tools) + except Exception as e: + renderer.print_warning(f"MCP load error: {e}") + + # ── Print banner ────────────────────────────────────── + renderer.print_banner( + model=session.model, + tools=connected_tools if connected_tools else None, + ) + + # ── Input session (history + autocomplete) ──────────── + history_dir = Path.home() / ".perplexity_cli" + history_dir.mkdir(parents=True, exist_ok=True) + + prompt_session = PromptSession( + history=FileHistory(str(history_dir / "input_history")), + auto_suggest=AutoSuggestFromHistory(), + style=PROMPT_STYLE, + mouse_support=False, + ) + + # ── Ctrl+C handler ──────────────────────────────────── + def handle_interrupt(sig, frame): + console.print("\n[dim]Use [bold]exit[/] to quit, or press Ctrl+D.[/]") + + signal.signal(signal.SIGINT, handle_interrupt) + + def get_tools_context() -> str: + tools = bridge.get_all_tools() + if not tools: + return "" + + lines = [ + "\n=== LOCAL SYSTEM ACCESS & MCP TOOLS ===", + "You have direct access to the local machine via the following Model Context Protocol (MCP) tools:", + ] + + for t in tools: + schema_str = json.dumps(t.get("input_schema", {}).get("properties", {})) + lines.append(f"- {t['full_name']}: {t['description']}") + lines.append(f" Schema: {schema_str}") + + lines.extend([ + "", + "CRITICAL DIRECTIVES FOR TOOL CALLS:", + "1. If the user asks you to perform local tasks (e.g., list files, check folder, read file, run shell command, check system stats), you MUST execute the appropriate tool instead of claiming you cannot access the system.", + "2. To invoke a tool, you MUST output exactly one line at the very start of your reply in this format:", + " TOOL_CALL: ", + " Example: TOOL_CALL: local.list_directory {\"path\": \".\"}", + " Example: TOOL_CALL: local.execute_command {\"command\": \"ls -la\"}", + "3. DO NOT search the web for local operations. Just call the tool directly.", + "4. Do NOT output any preamble, thoughts, markdown formatting, or conversational text. Output ONLY the TOOL_CALL line.", + "=========================================\n" + ]) + return "\n".join(lines) + + # ── Main REPL loop ──────────────────────────────────── + while True: + try: + user_input = prompt_session.prompt( + HTML('You '), + style=PROMPT_STYLE, + ).strip() + except KeyboardInterrupt: + continue + except EOFError: + console.print("\n[dim]Goodbye![/]") + break + + if not user_input: + continue + + # ── Slash commands ──────────────────────────────── + if user_input.startswith("/") or user_input.lower() in ("exit", "quit"): + result = handle_command(user_input, session, renderer, bridge) + if result == "exit": + console.print("[dim]Goodbye! 👋[/]") + bridge.disconnect_all() + break + continue + + # ── Send to Perplexity ──────────────────────────────────────── + session.add_user_message(user_input) + + try: + max_steps = 3 + step = 0 + while step < max_steps: + query_with_context = session.build_query_with_context(user_input) + + # Check if we are feeding back a tool result in the history + is_feedback = session.messages and session.messages[-1]["role"] == "tool" + + tools_context = get_tools_context() + if tools_context: + if is_feedback: + # Append an overriding system directive telling the LLM to synthesize the final answer + query_with_context += ( + "\n\n[SYSTEM DIRECTIVE: The requested tool has been successfully executed, " + "and the results are provided above in your Conversation History. " + "Do NOT output another TOOL_CALL directive. Use the results from the tool " + "to answer the user's original question directly now in plain, friendly text.]" + ) + else: + query_with_context += tools_context + + response = stream_perplexity( + client=client, + query=query_with_context, + model=session.model, + mode=session.mode, + search_focus=session.search_focus, + is_incognito=session.is_incognito, + renderer=renderer, + ) + + if not response.text: + break + + session.add_assistant_message(response.text, response.citations) + + # Execute tool call and check if one was triggered + tool_called = _maybe_handle_tool_call(response.text, bridge, session, renderer) + if not tool_called: + # Final response generated, break loop + break + + # Increment step and prepare to re-ask Perplexity + step += 1 + if step < max_steps: + console.print("\n[dim]Feeding tool results back to Perplexity...[/]\n") + + except KeyboardInterrupt: + console.print("\n[dim]Request cancelled.[/]\n") + except Exception as e: + renderer.print_error(f"Request failed: {e}") + + # ── Auto-save on exit ───────────────────────────────── + if session.message_count > 0: + try: + saved_path = session.save() + console.print(f"[dim]Session auto-saved to: {saved_path}[/]") + except Exception: + pass + + +def _maybe_handle_tool_call( + response_text: str, + bridge: MCPToolBridge, + session: ChatSession, + renderer: Renderer, +) -> bool: + """Auto-execute tool calls if LLM response contains TOOL_CALL: directives. + Uses highly robust splits and JSON substring balancing to handle duplicate, + concatenated, or conversation-polluted tool calls. + Returns True if at least one tool call was executed. + """ + if not bridge.has_tools: + return False + + import re + # Split by case-insensitive 'TOOL_CALL:' + parts = re.split(r'(?i)TOOL_CALL:', response_text) + if len(parts) <= 1: + return False + + called = False + seen_calls = set() + + for part in parts[1:]: + part = part.strip() + if not part: + continue + + # Split by first whitespace to separate tool name from arguments + sub_parts = part.split(None, 1) + tool_name = sub_parts[0].strip() + args_str = sub_parts[1].strip() if len(sub_parts) > 1 else "{}" + + # Highly robust JSON cleaning: if it starts with '{', find the longest valid JSON substring + if args_str.startswith("{"): + for idx in range(len(args_str), 0, -1): + sub = args_str[:idx] + if sub.endswith("}"): + try: + # Test if this substring is valid JSON + json.loads(sub) + args_str = sub + break + except json.JSONDecodeError: + pass + + # Deduplicate identical tool calls in the same turn + call_key = (tool_name, args_str) + if call_key in seen_calls: + continue + seen_calls.add(call_key) + + try: + args = json.loads(args_str) if args_str.startswith("{") else bridge.parse_tool_args(args_str) + renderer.print_tool_call(tool_name, args) + result, elapsed = bridge.call_tool(tool_name, args) + renderer.print_tool_result(tool_name, result, elapsed) + session.add_tool_result(tool_name, result) + called = True + except Exception as e: + renderer.print_tool_error(tool_name, str(e)) + called = True + + return called + + +# ─── Entry point ────────────────────────────────────────────────────────────── + +def main() -> None: + parser = argparse.ArgumentParser( + description="Perplexity Interactive CLI — multi-turn chat with MCP tool support", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python3 cli.py + python3 cli.py --model sonar-pro + python3 cli.py --model claude46sonnetthinking --mcp ./mcp_tools.json + python3 cli.py --load ~/.perplexity_cli/sessions/session_20240525.json + python3 cli.py --system "You are a cybersecurity expert." + python3 cli.py --incognito + """, + ) + parser.add_argument( + "--model", "-m", + default=os.getenv("DEFAULT_MODEL", DEFAULT_MODEL), + help=f"Model to use (default: {DEFAULT_MODEL})", + ) + parser.add_argument( + "--mcp", + metavar="CONFIG_JSON", + help="Path to MCP tools config JSON file", + ) + parser.add_argument( + "--load", + metavar="SESSION_JSON", + help="Load a saved session file", + ) + parser.add_argument( + "--mode", + default=os.getenv("DEFAULT_MODE", "copilot"), + choices=["copilot", "search", "reasoning", "writing"], + help="Search mode (default: copilot)", + ) + parser.add_argument( + "--focus", + default=os.getenv("DEFAULT_SEARCH_FOCUS", "internet"), + choices=["internet", "academic", "writing", "wolfram", "youtube", "reddit"], + help="Search focus (default: internet)", + ) + parser.add_argument( + "--system", + default=os.getenv("CLI_SYSTEM_PROMPT", ""), + help="System prompt to prepend to all queries", + ) + parser.add_argument( + "--incognito", + action="store_true", + help="Enable incognito mode", + ) + + args = parser.parse_args() + + run_cli( + model=args.model, + mcp_config=args.mcp, + load_session=args.load, + mode=args.mode, + search_focus=args.focus, + system_prompt=args.system or None, + incognito=args.incognito, + ) + + +if __name__ == "__main__": + main() diff --git a/mcp_tools.example.json b/mcp_tools.example.json new file mode 100644 index 0000000..7ea0099 --- /dev/null +++ b/mcp_tools.example.json @@ -0,0 +1,29 @@ +{ + "mcpServers": { + "burpsuite": { + "command": "python", + "args": ["/path/to/burpsuite-mcp/server.py"], + "env": { + "BURP_API_KEY": "your-burp-api-key", + "BURP_HOST": "http://127.0.0.1:1337" + } + }, + "filesystem": { + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/home/daud"], + "env": {} + }, + "brave-search": { + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-brave-search"], + "env": { + "BRAVE_API_KEY": "your-brave-api-key" + } + }, + "custom-tool": { + "command": "python", + "args": ["/path/to/your/custom_mcp_server.py"], + "env": {} + } + } +} diff --git a/mcp_tools.json b/mcp_tools.json new file mode 100644 index 0000000..df39a52 --- /dev/null +++ b/mcp_tools.json @@ -0,0 +1,8 @@ +{ + "mcpServers": { + "local": { + "command": "python3", + "args": ["/home/daud/AIProject/perplexity-scrape/scripts/shell_mcp.py"] + } + } +} diff --git a/requirements.txt b/requirements.txt index 9988ef1..97b9954 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,13 @@ curl_cffi>=0.7.0 python-dotenv>=1.0.0 -mcp[cli]>=1.0.0 +mcp[cli]>=1.9.0 fastapi>=0.115.0 uvicorn[standard]>=0.30.0 pydantic>=2.0.0 pytest>=8.0.0 pytest-asyncio>=0.23.0 httpx>=0.27.0 + +# CLI dependencies +rich>=13.7.0 +prompt_toolkit>=3.0.43 diff --git a/scripts/shell_mcp.py b/scripts/shell_mcp.py new file mode 100644 index 0000000..ffee9e6 --- /dev/null +++ b/scripts/shell_mcp.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +""" +Local Shell & Filesystem MCP Server + +Provides standard shell execution and filesystem tools to the Perplexity CLI agentic loop. +""" + +import os +import subprocess +from mcp.server.fastmcp import FastMCP + +mcp = FastMCP("Local Shell & Filesystem") + + +@mcp.tool() +def list_directory(path: str = ".") -> str: + """List contents of a directory. + + Args: + path: Path to the directory (default: "."). + """ + try: + resolved = os.path.abspath(path) + items = os.listdir(resolved) + lines = [f"Directory: {resolved}", ""] + for item in sorted(items): + full_path = os.path.join(resolved, item) + suffix = "/" if os.path.isdir(full_path) else "" + size = f" ({os.path.getsize(full_path)} bytes)" if os.path.isfile(full_path) else "" + lines.append(f"- {item}{suffix}{size}") + return "\n".join(lines) + except Exception as e: + return f"Error: {e}" + + +@mcp.tool() +def read_file(path: str) -> str: + """Read the contents of a local file. + + Args: + path: Path to the file. + """ + try: + resolved = os.path.abspath(path) + if not os.path.isfile(resolved): + return f"Error: Path is not a file: {resolved}" + with open(resolved, "r", encoding="utf-8", errors="replace") as f: + return f.read() + except Exception as e: + return f"Error: {e}" + + +@mcp.tool() +def execute_command(command: str, cwd: str = ".") -> str: + """Execute a bash shell command locally and return stdout/stderr. + + Args: + command: The shell command to run (e.g. 'uname -a', 'python3 -m pip list'). + cwd: Working directory to run the command in (default: "."). + """ + try: + resolved_cwd = os.path.abspath(cwd) + # Execute with 30s timeout to prevent hanging the CLI loop + result = subprocess.run( + command, + shell=True, + cwd=resolved_cwd, + text=True, + capture_output=True, + timeout=30, + ) + output = [] + if result.stdout: + output.append("[Stdout]") + output.append(result.stdout) + if result.stderr: + output.append("[Stderr]") + output.append(result.stderr) + if not output: + output.append(f"[Process exited with code {result.returncode}]") + return "\n".join(output) + except subprocess.TimeoutExpired: + return "Error: Command timed out after 30 seconds." + except Exception as e: + return f"Error executing command: {e}" + + +if __name__ == "__main__": + mcp.run() diff --git a/src/cli/__init__.py b/src/cli/__init__.py new file mode 100644 index 0000000..ad08ba7 --- /dev/null +++ b/src/cli/__init__.py @@ -0,0 +1,9 @@ +""" +CLI module for Perplexity interactive shell. +""" + +from .session import ChatSession +from .renderer import Renderer +from .mcp_bridge import MCPToolBridge + +__all__ = ["ChatSession", "Renderer", "MCPToolBridge"] diff --git a/src/cli/mcp_bridge.py b/src/cli/mcp_bridge.py new file mode 100644 index 0000000..4b09909 --- /dev/null +++ b/src/cli/mcp_bridge.py @@ -0,0 +1,322 @@ +""" +MCP Tool Bridge + +Connects to external MCP servers (BurpSuite, filesystem, custom, etc.) +and exposes their tools to the CLI's agentic loop. + +Supports stdio-based MCP servers loaded from a JSON config file. +""" + +import asyncio +import json +import subprocess +import time +from pathlib import Path +from typing import Any, Optional + +try: + from mcp import ClientSession, StdioServerParameters + from mcp.client.stdio import stdio_client + HAS_MCP = True +except ImportError: + HAS_MCP = False + + +class MCPServer: + """Represents a connected MCP server with its available tools.""" + + def __init__(self, name: str, config: dict): + self.name = name + self.config = config + self.tools: list[dict] = [] + self.session: Any = None + self._client_ctx: Any = None + self._session_ctx: Any = None + self.connected = False + + async def connect(self) -> bool: + """Connect to the MCP server via stdio.""" + if not HAS_MCP: + return False + + try: + command = self.config.get("command", "") + args = self.config.get("args", []) + env = self.config.get("env", None) + + server_params = StdioServerParameters( + command=command, + args=args, + env=env, + ) + + self._client_ctx = stdio_client(server_params) + read, write = await self._client_ctx.__aenter__() + + self._session_ctx = ClientSession(read, write) + self.session = await self._session_ctx.__aenter__() + + await self.session.initialize() + + # Fetch available tools + tools_result = await self.session.list_tools() + self.tools = [ + { + "name": t.name, + "server": self.name, + "full_name": f"{self.name}.{t.name}", + "description": t.description or "", + "input_schema": t.inputSchema or {}, + } + for t in tools_result.tools + ] + self.connected = True + return True + + except Exception as e: + self.connected = False + raise ConnectionError(f"Failed to connect to MCP server '{self.name}': {e}") + + async def call_tool(self, tool_name: str, arguments: dict) -> str: + """Execute a tool on this MCP server.""" + if not self.session: + raise RuntimeError(f"Not connected to server '{self.name}'") + + result = await self.session.call_tool(tool_name, arguments) + + # Extract text content from result + if result.content: + parts = [] + for item in result.content: + if hasattr(item, "text"): + parts.append(item.text) + elif hasattr(item, "data"): + parts.append(f"[binary data: {len(item.data)} bytes]") + else: + parts.append(str(item)) + return "\n".join(parts) + + return "(empty result)" + + async def disconnect(self) -> None: + """Disconnect from the MCP server.""" + try: + if self._session_ctx: + await self._session_ctx.__aexit__(None, None, None) + if self._client_ctx: + await self._client_ctx.__aexit__(None, None, None) + except Exception: + pass + self.connected = False + + +class MCPToolBridge: + """ + Manages connections to multiple MCP servers and provides + a unified interface for listing and calling tools. + """ + + def __init__(self): + self.servers: dict[str, MCPServer] = {} + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._thread: Optional[threading.Thread] = None + + def _start_loop(self) -> None: + """Start the background event loop in a daemon thread if not already running.""" + import threading + if self._loop is None or self._loop.is_closed(): + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._run_loop, daemon=True) + self._thread.start() + + def _run_loop(self) -> None: + """Background thread target to run the event loop forever.""" + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def load_config(self, config_path: str) -> list[str]: + """ + Load MCP server configs from a JSON file. + Returns list of server names that were loaded. + """ + path = Path(config_path) + if not path.exists(): + raise FileNotFoundError(f"Config file not found: {config_path}") + + with open(path) as f: + config = json.load(f) + + servers_config = config.get("mcpServers", {}) + loaded = [] + + for name, server_cfg in servers_config.items(): + self.servers[name] = MCPServer(name, server_cfg) + loaded.append(name) + + return loaded + + def connect_all(self) -> dict[str, bool]: + """Connect to all configured servers thread-safely. Returns {name: success} map.""" + if not self.servers: + return {} + + self._start_loop() + results = {} + + for name, server in self.servers.items(): + try: + # Dispatch connect asynchronously to background thread loop + future = asyncio.run_coroutine_threadsafe(server.connect(), self._loop) + # Wait up to 10 seconds for connection to succeed + success = future.result(timeout=10) + results[name] = success + except Exception as e: + results[name] = False + server.last_error = str(e) + + return results + + def connect_server(self, name: str) -> bool: + """Connect to a specific server by name thread-safely.""" + if name not in self.servers: + raise KeyError(f"No server named '{name}'") + + self._start_loop() + future = asyncio.run_coroutine_threadsafe(self.servers[name].connect(), self._loop) + return future.result(timeout=10) + + def get_all_tools(self) -> list[dict]: + """Return all tools from all connected servers.""" + tools = [] + for server in self.servers.values(): + if server.connected: + tools.extend(server.tools) + return tools + + def find_tool(self, full_name: str) -> Optional[tuple[MCPServer, str]]: + """ + Find a server and tool name from a full_name like 'burpsuite.scan'. + Returns (server, tool_name) or None. + """ + if "." in full_name: + server_name, tool_name = full_name.split(".", 1) + if server_name in self.servers and self.servers[server_name].connected: + return self.servers[server_name], tool_name + else: + # Search by tool name across all servers + for server in self.servers.values(): + if server.connected: + for tool in server.tools: + if tool["name"] == full_name: + return server, full_name + return None + + def call_tool(self, full_name: str, arguments: dict) -> tuple[str, float]: + """ + Call a tool by full name (e.g. 'burpsuite.active_scan') thread-safely. + Returns (result, elapsed_seconds). + """ + found = self.find_tool(full_name) + if not found: + raise KeyError(f"Tool '{full_name}' not found in any connected server") + + server, tool_name = found + self._start_loop() + + start = time.time() + # Dispatch execution asynchronously to background thread loop + future = asyncio.run_coroutine_threadsafe( + server.call_tool(tool_name, arguments), self._loop + ) + try: + # Force 30 second timeout on tool execution to prevent freezes + result = future.result(timeout=30) + except Exception as e: + result = f"Error: Tool call timed out or failed: {e}" + + elapsed = time.time() - start + return result, elapsed + + def parse_tool_args(self, args_str: str) -> dict: + """ + Parse tool arguments from CLI string. + Supports: key=value pairs OR raw JSON string. + + Examples: + url=https://example.com scan_type=full + {"url": "https://example.com", "scan_type": "full"} + """ + args_str = args_str.strip() + if not args_str: + return {} + + # Try JSON first + if args_str.startswith("{"): + try: + return json.loads(args_str) + except json.JSONDecodeError: + pass + + # Parse key=value pairs + result = {} + parts = args_str.split() + for part in parts: + if "=" in part: + key, _, value = part.partition("=") + # Try to cast to int/float/bool + if value.lower() == "true": + result[key] = True + elif value.lower() == "false": + result[key] = False + else: + try: + result[key] = int(value) + except ValueError: + try: + result[key] = float(value) + except ValueError: + result[key] = value + else: + # Positional arg + result[f"arg_{len(result)}"] = part + + return result + + def get_tools_as_openai_format(self) -> list[dict]: + """ + Return all tools in OpenAI function-calling format. + Useful for injecting into LLM context. + """ + tools = [] + for tool in self.get_all_tools(): + tools.append({ + "type": "function", + "function": { + "name": tool["full_name"].replace(".", "__"), + "description": tool["description"], + "parameters": tool.get("input_schema", {"type": "object", "properties": {}}), + }, + }) + return tools + + def disconnect_all(self) -> None: + """Disconnect from all servers thread-safely and shut down loop.""" + if self._loop and not self._loop.is_closed(): + for server in self.servers.values(): + try: + future = asyncio.run_coroutine_threadsafe(server.disconnect(), self._loop) + future.result(timeout=3) + except Exception: + pass + try: + self._loop.call_soon_threadsafe(self._loop.stop) + except Exception: + pass + + @property + def has_tools(self) -> bool: + return bool(self.get_all_tools()) + + @property + def connected_servers(self) -> list[str]: + return [n for n, s in self.servers.items() if s.connected] diff --git a/src/cli/renderer.py b/src/cli/renderer.py new file mode 100644 index 0000000..d5582c1 --- /dev/null +++ b/src/cli/renderer.py @@ -0,0 +1,353 @@ +""" +Terminal Renderer + +Rich-based output renderer for the Perplexity CLI. +Handles markdown, streaming tokens, citations, and tool call display. +""" + +import sys +from typing import Optional +from datetime import datetime + +from rich.console import Console +from rich.markdown import Markdown +from rich.panel import Panel +from rich.table import Table +from rich.text import Text +from rich.rule import Rule +from rich.live import Live +from rich.spinner import Spinner +from rich.columns import Columns +from rich import box + + +console = Console() +error_console = Console(stderr=True, style="bold red") + + +class Renderer: + """Handles all terminal output formatting.""" + + BRAND_COLOR = "bright_cyan" + USER_COLOR = "bright_green" + ASSISTANT_COLOR = "bright_white" + TOOL_COLOR = "bright_yellow" + DIM_COLOR = "dim" + ERROR_COLOR = "bright_red" + CITATION_COLOR = "cyan" + + def __init__(self, no_color: bool = False, compact: bool = False): + self.no_color = no_color + self.compact = compact + + def print_banner(self, model: str, tools: list[str] | None = None) -> None: + """Print the startup banner.""" + tool_str = ", ".join(tools) if tools else "none" + + banner = Text() + banner.append(" Perplexity", style="bold bright_cyan") + banner.append(" CLI ", style="bold white") + banner.append("⚡", style="bright_yellow") + + subtitle = Text() + subtitle.append(" Model: ", style="dim") + subtitle.append(model, style="bold bright_cyan") + subtitle.append(" Tools: ", style="dim") + subtitle.append(tool_str, style="bright_yellow" if tools else "dim") + + console.print() + console.print(Panel( + f"{banner}\n{subtitle}", + border_style="bright_cyan", + padding=(0, 1), + )) + console.print( + " Type [bold]/help[/bold] for commands · [bold]exit[/bold] to quit", + style="dim" + ) + console.print() + + def print_help(self) -> None: + """Print the help table.""" + table = Table( + title="Available Commands", + box=box.ROUNDED, + border_style="bright_cyan", + show_header=True, + header_style="bold bright_cyan", + padding=(0, 1), + ) + table.add_column("Command", style="bold bright_yellow", no_wrap=True) + table.add_column("Description", style="white") + table.add_column("Example", style="dim") + + commands = [ + ("/model ", "Switch active model", "/model sonar-pro"), + ("/models", "List all available models", "/models"), + ("/clear", "Clear conversation history", "/clear"), + ("/save [path]", "Save session to JSON file", "/save ~/chat.json"), + ("/load ", "Load a saved session", "/load ~/chat.json"), + ("/sessions", "List saved sessions", "/sessions"), + ("/system ", "Set system prompt", "/system You are a security expert"), + ("/system", "Show current system prompt", "/system"), + ("/mode ", "Set search mode (copilot/search)", "/mode search"), + ("/focus ", "Set search focus", "/focus academic"), + ("/incognito", "Toggle incognito mode", "/incognito"), + ("/tools", "List connected MCP tools", "/tools"), + ("/tool [a]", "Call MCP tool manually", "/tool burpsuite.scan url=..."), + ("/mcp ", "Load MCP servers from JSON file", "/mcp ./mcp_tools.json"), + ("/history", "Show conversation history", "/history"), + ("/status", "Show session status", "/status"), + ("/help", "Show this help", "/help"), + ("exit / quit", "Exit the CLI", "exit"), + ] + + for cmd, desc, example in commands: + table.add_row(cmd, desc, example) + + console.print(table) + console.print() + + def print_status(self, session) -> None: + """Print current session status.""" + table = Table(box=box.SIMPLE, show_header=False, padding=(0, 1)) + table.add_column("Key", style="dim") + table.add_column("Value", style="bright_white") + + table.add_row("Model", f"[bold bright_cyan]{session.model}[/]") + table.add_row("Mode", session.mode) + table.add_row("Focus", session.search_focus) + table.add_row("Messages", str(session.message_count)) + table.add_row("Incognito", "✓ On" if session.is_incognito else "✗ Off") + table.add_row("Session ID", session.session_id) + if session.system_prompt: + prompt_preview = session.system_prompt[:60] + "…" if len(session.system_prompt) > 60 else session.system_prompt + table.add_row("System", prompt_preview) + else: + table.add_row("System", "[dim]not set[/]") + + console.print(Panel(table, title="Session Status", border_style="bright_cyan")) + console.print() + + def print_user_prompt(self) -> str: + """Print user prompt prefix (for manual input loops).""" + return "" + + def print_user_message(self, content: str) -> None: + """Echo the user message with styling.""" + console.print() + console.print(f"[bold {self.USER_COLOR}]You[/] [dim]▸[/] {content}") + console.print() + + def start_assistant_stream(self, model: str) -> Live: + """Start a live spinner while waiting for first token.""" + spinner = Spinner("dots2", text=f" [dim]Thinking ({model})…[/]", style="bright_cyan") + return Live(spinner, console=console, refresh_per_second=10, transient=True) + + def print_assistant_response( + self, + text: str, + citations: list[dict] | None = None, + model: str = "", + elapsed: float = 0, + ) -> None: + """Render the full assistant response with markdown and citations.""" + # Header + header = Text() + header.append("Perplexity", style=f"bold {self.BRAND_COLOR}") + if model: + header.append(f" [{model}]", style="dim") + if elapsed: + header.append(f" {elapsed:.1f}s", style="dim") + + console.print(header) + console.print(Rule(style="dim cyan")) + + # Main response (markdown rendered) + md = Markdown(text, code_theme="monokai") + console.print(md) + + # Citations + if citations: + self._print_citations(citations) + + console.print() + + def _print_citations(self, citations: list[dict]) -> None: + """Print citations as a compact footnote list.""" + console.print() + console.print("[dim]── Sources ──────────────────────────────────────[/]") + for i, cite in enumerate(citations[:8], 1): # limit to 8 + title = cite.get("title", "Unknown") + url = cite.get("url", "") + console.print( + f" [{self.CITATION_COLOR}][{i}][/] [link={url}]{title}[/link]", + style="dim" + ) + + def print_tool_call(self, tool_name: str, args: dict) -> None: + """Print a tool call being executed.""" + args_str = " ".join(f"[dim]{k}=[/][bright_white]{v}[/]" for k, v in args.items()) + console.print( + f" [bright_yellow]🔧 Tool[/] [bold]{tool_name}[/] {args_str}" + ) + + def print_tool_result(self, tool_name: str, result: str, elapsed: float = 0) -> None: + """Print a tool result.""" + elapsed_str = f" ({elapsed:.1f}s)" if elapsed else "" + preview = result[:200] + "…" if len(result) > 200 else result + console.print( + f" [green]✓ Result[/] [dim]{tool_name}{elapsed_str}[/]" + ) + if preview: + console.print(f" [dim]{preview}[/]") + console.print() + + def print_tool_error(self, tool_name: str, error: str) -> None: + """Print a tool execution error.""" + console.print( + f" [bright_red]✗ Tool Error[/] [bold]{tool_name}[/]: {error}" + ) + + def print_models_table(self, models: list, current: str) -> None: + """Print available models. Accepts list of str IDs or dicts {id, name, thinking}.""" + table = Table( + title="Available Models", + box=box.ROUNDED, + border_style="bright_cyan", + show_header=True, + header_style="bold bright_cyan", + ) + table.add_column("#", style="dim", width=3) + table.add_column("ID", style="bold", no_wrap=True) + table.add_column("Name", style="white") + table.add_column("", justify="center", width=10) # thinking badge + table.add_column("Active", justify="center", width=6) + + for i, m in enumerate(models, 1): + if isinstance(m, dict): + mid = m["id"] + mname = m.get("name", mid) + thinking = m.get("thinking", False) + else: + mid = m + mname = "" + thinking = False + + is_active = mid == current + active_mark = "[bold bright_green]●[/]" if is_active else "[dim]○[/]" + id_style = "bold bright_cyan" if is_active else "white" + think_badge = "[bright_yellow]⚡ thinking[/]" if thinking else "" + + table.add_row(str(i), f"[{id_style}]{mid}[/]", mname, think_badge, active_mark) + + console.print(table) + console.print() + + def print_tools_table(self, tools: list[dict]) -> None: + """Print connected MCP tools.""" + if not tools: + console.print("[dim]No MCP tools connected. Use /mcp to load.[/]") + console.print() + return + + table = Table( + title="Connected MCP Tools", + box=box.ROUNDED, + border_style="bright_yellow", + show_header=True, + header_style="bold bright_yellow", + ) + table.add_column("Tool Name", style="bold bright_yellow") + table.add_column("Server", style="dim") + table.add_column("Description", style="white") + + for tool in tools: + table.add_row( + tool.get("name", "?"), + tool.get("server", "?"), + tool.get("description", "")[:60], + ) + + console.print(table) + console.print() + + def print_history(self, messages: list[dict]) -> None: + """Print conversation history.""" + if not messages: + console.print("[dim]No conversation history yet.[/]") + return + + console.print(Rule("Conversation History", style="bright_cyan")) + for msg in messages: + role = msg["role"] + content = msg.get("content", "") + ts = msg.get("timestamp", "")[:19] if msg.get("timestamp") else "" + + if role == "user": + console.print(f"\n[bold bright_green]You[/] [dim]{ts}[/]") + console.print(f" {content}") + elif role == "assistant": + console.print(f"\n[bold bright_cyan]Perplexity[/] [dim]{ts}[/]") + md = Markdown(content[:500] + ("…" if len(content) > 500 else "")) + console.print(md) + elif role == "tool": + console.print( + f"\n[bold bright_yellow]Tool: {msg.get('tool_name', '?')}[/] [dim]{ts}[/]" + ) + console.print(f" [dim]{content[:200]}[/]") + + console.print(Rule(style="bright_cyan")) + console.print() + + def print_sessions_table(self, sessions: list[dict]) -> None: + """Print saved sessions.""" + if not sessions: + console.print("[dim]No saved sessions found.[/]") + console.print() + return + + table = Table( + title="Saved Sessions", + box=box.ROUNDED, + border_style="bright_cyan", + show_header=True, + header_style="bold bright_cyan", + ) + table.add_column("Session ID", style="bold") + table.add_column("Created", style="dim") + table.add_column("Model", style="bright_cyan") + table.add_column("Messages", justify="right") + table.add_column("Path", style="dim") + + for s in sessions: + table.add_row( + s["session_id"], + s["created_at"][:19], + s["model"], + str(s["message_count"]), + s["path"], + ) + + console.print(table) + console.print() + + def print_info(self, msg: str) -> None: + console.print(f"[bright_cyan]ℹ[/] {msg}") + + def print_success(self, msg: str) -> None: + console.print(f"[bright_green]✓[/] {msg}") + + def print_warning(self, msg: str) -> None: + console.print(f"[bright_yellow]⚠[/] {msg}") + + def print_error(self, msg: str) -> None: + console.print(f"[bright_red]✗[/] {msg}") + + def print_divider(self) -> None: + console.print(Rule(style="dim")) + + def print_cleared(self) -> None: + console.clear() + console.print("[bright_green]✓[/] Conversation history cleared.") + console.print() diff --git a/src/cli/session.py b/src/cli/session.py new file mode 100644 index 0000000..6e22258 --- /dev/null +++ b/src/cli/session.py @@ -0,0 +1,267 @@ +""" +Chat Session Manager + +Manages conversation history, model selection, and session persistence. +Models are loaded dynamically from opencode-provider.json. +""" + +import json +import os +from datetime import datetime +from pathlib import Path +from typing import Optional + + +# ─── Dynamic model loading from opencode-provider.json ─────────────────────── + +def _load_models_from_provider(provider_path: Optional[str] = None) -> list[dict]: + """ + Load model list from opencode-provider.json. + Returns list of {id, name, thinking, context} dicts. + """ + search_paths = [ + provider_path, + Path(__file__).parent.parent.parent / "opencode-provider.json", + Path.cwd() / "opencode-provider.json", + ] + + for path in search_paths: + if path is None: + continue + p = Path(path) + if p.exists(): + try: + with open(p) as f: + data = json.load(f) + # Structure: {"perplexity-scrape": {"models": {id: {name, ...}}}} + for provider_data in data.values(): + models_dict = provider_data.get("models", {}) + result = [] + for model_id, model_info in models_dict.items(): + result.append({ + "id": model_id, + "name": model_info.get("name", model_id), + "thinking": model_info.get("thinking", False), + "context": model_info.get("limit", {}).get("context", 0), + }) + if result: + return result + except Exception: + pass + + # Fallback hardcoded list (mirrors opencode-provider.json) + return [ + {"id": "claude45sonnetthinking", "name": "Claude 4.5 Sonnet + Reasoning", "thinking": True, "context": 200000}, + {"id": "claude45sonnet", "name": "Claude 4.5 Sonnet", "thinking": False, "context": 200000}, + {"id": "claude45opusthinking", "name": "Claude 4.5 Opus + Reasoning", "thinking": True, "context": 200000}, + {"id": "claude45opus", "name": "Claude 4.5 Opus", "thinking": False, "context": 200000}, + {"id": "gemini30flash", "name": "Gemini 3 Flash", "thinking": False, "context": 1048576}, + {"id": "gemini30flash_high", "name": "Gemini 3 Flash + Reasoning", "thinking": True, "context": 1048576}, + {"id": "gemini30pro", "name": "Gemini 3 Pro + Reasoning", "thinking": True, "context": 1048576}, + {"id": "gpt52", "name": "GPT-5.2", "thinking": False, "context": 128000}, + {"id": "gpt52_thinking", "name": "GPT-5.2 + Reasoning", "thinking": True, "context": 128000}, + {"id": "grok41nonreasoning", "name": "Grok 4.1", "thinking": False, "context": 128000}, + {"id": "grok41reasoning", "name": "Grok 4.1 + Reasoning", "thinking": True, "context": 128000}, + {"id": "kimik25thinking", "name": "Kimi K2.5 Thinking", "thinking": True, "context": 128000}, + {"id": "sonar", "name": "Sonar (Experimental)", "thinking": False, "context": 128000}, + ] + + +# Load at import time +_ALL_MODELS: list[dict] = _load_models_from_provider() +AVAILABLE_MODELS: list[str] = [m["id"] for m in _ALL_MODELS] +DEFAULT_MODEL: str = AVAILABLE_MODELS[0] if AVAILABLE_MODELS else "claude45sonnetthinking" + + +def get_model_info(model_id: str) -> dict: + """Get full info dict for a model ID.""" + for m in _ALL_MODELS: + if m["id"] == model_id: + return m + return {"id": model_id, "name": model_id, "thinking": False, "context": 0} + + +def get_model_name(model_id: str) -> str: + """Get display name for a model ID.""" + return get_model_info(model_id).get("name", model_id) + + +def reload_models(provider_path: Optional[str] = None) -> list[str]: + """Reload model list from opencode-provider.json.""" + global _ALL_MODELS, AVAILABLE_MODELS, DEFAULT_MODEL + _ALL_MODELS = _load_models_from_provider(provider_path) + AVAILABLE_MODELS = [m["id"] for m in _ALL_MODELS] + DEFAULT_MODEL = AVAILABLE_MODELS[0] if AVAILABLE_MODELS else "claude45sonnetthinking" + return AVAILABLE_MODELS + + +MAX_HISTORY_MESSAGES = 50 # keep last N messages to avoid token overflow + + +class ChatSession: + """Manages a single chat session with history and model config.""" + + def __init__( + self, + model: str = DEFAULT_MODEL, + mode: str = "copilot", + search_focus: str = "internet", + system_prompt: Optional[str] = None, + is_incognito: bool = False, + ): + self.model = model + self.mode = mode + self.search_focus = search_focus + self.system_prompt = system_prompt + self.is_incognito = is_incognito + self.messages: list[dict] = [] + self.created_at = datetime.now() + self.session_id = datetime.now().strftime("%Y%m%d_%H%M%S") + + def add_user_message(self, content: str) -> None: + """Add a user message to history.""" + self.messages.append({ + "role": "user", + "content": content, + "timestamp": datetime.now().isoformat(), + }) + + def add_assistant_message(self, content: str, citations: list[dict] | None = None) -> None: + """Add an assistant message to history.""" + msg: dict = { + "role": "assistant", + "content": content, + "timestamp": datetime.now().isoformat(), + } + if citations: + msg["citations"] = citations + self.messages.append(msg) + + def add_tool_result(self, tool_name: str, result: str) -> None: + """Add a tool execution result to history.""" + self.messages.append({ + "role": "tool", + "tool_name": tool_name, + "content": result, + "timestamp": datetime.now().isoformat(), + }) + + def get_history_for_display(self) -> list[dict]: + """Return messages formatted for display.""" + return self.messages + + def build_query_with_context(self, new_query: str) -> str: + """ + Build the full query string including conversation history context. + Perplexity doesn't have native multi-turn, so we inject history as context. + """ + if not self.messages: + return new_query + + recent = self.messages[-MAX_HISTORY_MESSAGES:] + context_parts = [] + + if self.system_prompt: + context_parts.append(f"[System Instructions]\n{self.system_prompt}\n") + + context_parts.append("[Conversation History]") + for msg in recent: + role = msg["role"].upper() + content = msg["content"] + if role == "TOOL": + context_parts.append(f"[Tool: {msg.get('tool_name', '?')}]\n{content}") + else: + context_parts.append(f"{role}: {content}") + + context_parts.append(f"\n[Current Question]\n{new_query}") + context_parts.append( + "\nPlease answer the current question, taking the conversation history into account." + ) + + return "\n\n".join(context_parts) + + def clear(self) -> None: + """Clear conversation history.""" + self.messages = [] + + def set_model(self, model: str) -> bool: + """Change the active model. Returns True if valid.""" + self.model = model + return True + + def set_system_prompt(self, prompt: str) -> None: + """Set or update the system prompt.""" + self.system_prompt = prompt + + def save(self, path: Optional[str] = None) -> str: + """Save session to JSON file. Returns the file path.""" + if path is None: + save_dir = Path.home() / ".perplexity_cli" / "sessions" + save_dir.mkdir(parents=True, exist_ok=True) + path = str(save_dir / f"session_{self.session_id}.json") + + data = { + "session_id": self.session_id, + "created_at": self.created_at.isoformat(), + "model": self.model, + "mode": self.mode, + "search_focus": self.search_focus, + "system_prompt": self.system_prompt, + "is_incognito": self.is_incognito, + "messages": self.messages, + } + + with open(path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + + return path + + @classmethod + def load(cls, path: str) -> "ChatSession": + """Load a session from a JSON file.""" + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + + session = cls( + model=data.get("model", DEFAULT_MODEL), + mode=data.get("mode", "copilot"), + search_focus=data.get("search_focus", "internet"), + system_prompt=data.get("system_prompt"), + is_incognito=data.get("is_incognito", False), + ) + session.session_id = data.get("session_id", session.session_id) + session.messages = data.get("messages", []) + + return session + + @staticmethod + def list_saved_sessions() -> list[dict]: + """List all saved sessions.""" + save_dir = Path.home() / ".perplexity_cli" / "sessions" + if not save_dir.exists(): + return [] + + sessions = [] + for f in sorted(save_dir.glob("session_*.json"), reverse=True): + try: + with open(f) as fp: + data = json.load(fp) + sessions.append({ + "path": str(f), + "session_id": data.get("session_id", "?"), + "created_at": data.get("created_at", "?"), + "model": data.get("model", "?"), + "message_count": len(data.get("messages", [])), + }) + except Exception: + pass + + return sessions + + @property + def message_count(self) -> int: + return len(self.messages) + + @property + def user_messages(self) -> int: + return sum(1 for m in self.messages if m["role"] == "user")